IQ.Lvbs/IQ.Pilot
8
7
Fork 3
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

IQ.Pilot Release Commit @ f2a861c

Browse Source
This commit is contained in:
IQ.Lvbs CI [bot]
2026-09-02 15:07:09 -05:00
parent b42569dbca
commit e8748fd704
5497 changed files with 316070 additions and 179848 deletions
Download Patch File Download Diff File Expand all files Collapse all files

416
artifacts/package_sources/tinygrad/extra/amdpci/am_smi.py Executable file
View File

@@ -0,0 +1,416 @@
#!/usr/bin/env python3
import time, mmap, sys, shutil, os, glob, subprocess, argparse, collections
from tinygrad.helpers import DEBUG, NO_COLOR, 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 s if NO_COLOR else 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(fallback=(231, 24))
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...")

279
artifacts/package_sources/tinygrad/extra/amdpci/headers/amdgpu_doorbell.h Normal file
View File

@@ -0,0 +1,279 @@
/*
* Copyright 2018 Advanced Micro Devices, Inc.
*
* Permission is hereby granted, free of charge, to any person obtaining a
* copy of this software and associated documentation files (the "Software"),
* to deal in the Software without restriction, including without limitation
* the rights to use, copy, modify, merge, publish, distribute, sublicense,
* and/or sell copies of the Software, and to permit persons to whom the
* Software is furnished to do so, subject to the following conditions:
*
* The above copyright notice and this permission notice shall be included in
* all copies or substantial portions of the Software.
*
* THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
* IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
* FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
* THE COPYRIGHT HOLDER(S) OR AUTHOR(S) BE LIABLE FOR ANY CLAIM, DAMAGES OR
* OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE,
* ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
* OTHER DEALINGS IN THE SOFTWARE.
*
*/
#ifndef AMDGPU_DOORBELL_H
#define AMDGPU_DOORBELL_H
enum AMDGPU_DOORBELL_ASSIGNMENT {
AMDGPU_DOORBELL_KIQ = 0x000,
AMDGPU_DOORBELL_HIQ = 0x001,
AMDGPU_DOORBELL_DIQ = 0x002,
AMDGPU_DOORBELL_MEC_RING0 = 0x010,
AMDGPU_DOORBELL_MEC_RING1 = 0x011,
AMDGPU_DOORBELL_MEC_RING2 = 0x012,
AMDGPU_DOORBELL_MEC_RING3 = 0x013,
AMDGPU_DOORBELL_MEC_RING4 = 0x014,
AMDGPU_DOORBELL_MEC_RING5 = 0x015,
AMDGPU_DOORBELL_MEC_RING6 = 0x016,
AMDGPU_DOORBELL_MEC_RING7 = 0x017,
AMDGPU_DOORBELL_GFX_RING0 = 0x020,
AMDGPU_DOORBELL_sDMA_ENGINE0 = 0x1E0,
AMDGPU_DOORBELL_sDMA_ENGINE1 = 0x1E1,
AMDGPU_DOORBELL_IH = 0x1E8,
AMDGPU_DOORBELL_MAX_ASSIGNMENT = 0x3FF,
AMDGPU_DOORBELL_INVALID = 0xFFFF
};
enum AMDGPU_VEGA20_DOORBELL_ASSIGNMENT {
/* Compute + GFX: 0~255 */
AMDGPU_VEGA20_DOORBELL_KIQ = 0x000,
AMDGPU_VEGA20_DOORBELL_HIQ = 0x001,
AMDGPU_VEGA20_DOORBELL_DIQ = 0x002,
AMDGPU_VEGA20_DOORBELL_MEC_RING0 = 0x003,
AMDGPU_VEGA20_DOORBELL_MEC_RING1 = 0x004,
AMDGPU_VEGA20_DOORBELL_MEC_RING2 = 0x005,
AMDGPU_VEGA20_DOORBELL_MEC_RING3 = 0x006,
AMDGPU_VEGA20_DOORBELL_MEC_RING4 = 0x007,
AMDGPU_VEGA20_DOORBELL_MEC_RING5 = 0x008,
AMDGPU_VEGA20_DOORBELL_MEC_RING6 = 0x009,
AMDGPU_VEGA20_DOORBELL_MEC_RING7 = 0x00A,
AMDGPU_VEGA20_DOORBELL_USERQUEUE_START = 0x00B,
AMDGPU_VEGA20_DOORBELL_USERQUEUE_END = 0x08A,
AMDGPU_VEGA20_DOORBELL_GFX_RING0 = 0x08B,
/* SDMA:256~335*/
AMDGPU_VEGA20_DOORBELL_sDMA_ENGINE0 = 0x100,
AMDGPU_VEGA20_DOORBELL_sDMA_ENGINE1 = 0x10A,
AMDGPU_VEGA20_DOORBELL_sDMA_ENGINE2 = 0x114,
AMDGPU_VEGA20_DOORBELL_sDMA_ENGINE3 = 0x11E,
AMDGPU_VEGA20_DOORBELL_sDMA_ENGINE4 = 0x128,
AMDGPU_VEGA20_DOORBELL_sDMA_ENGINE5 = 0x132,
AMDGPU_VEGA20_DOORBELL_sDMA_ENGINE6 = 0x13C,
AMDGPU_VEGA20_DOORBELL_sDMA_ENGINE7 = 0x146,
/* IH: 376~391 */
AMDGPU_VEGA20_DOORBELL_IH = 0x178,
/* MMSCH: 392~407
* overlap the doorbell assignment with VCN as they are mutually exclusive
* VCN engine's doorbell is 32 bit and two VCN ring share one QWORD
*/
AMDGPU_VEGA20_DOORBELL64_VCN0_1 = 0x188, /* VNC0 */
AMDGPU_VEGA20_DOORBELL64_VCN2_3 = 0x189,
AMDGPU_VEGA20_DOORBELL64_VCN4_5 = 0x18A,
AMDGPU_VEGA20_DOORBELL64_VCN6_7 = 0x18B,
AMDGPU_VEGA20_DOORBELL64_VCN8_9 = 0x18C, /* VNC1 */
AMDGPU_VEGA20_DOORBELL64_VCNa_b = 0x18D,
AMDGPU_VEGA20_DOORBELL64_VCNc_d = 0x18E,
AMDGPU_VEGA20_DOORBELL64_VCNe_f = 0x18F,
AMDGPU_VEGA20_DOORBELL64_UVD_RING0_1 = 0x188,
AMDGPU_VEGA20_DOORBELL64_UVD_RING2_3 = 0x189,
AMDGPU_VEGA20_DOORBELL64_UVD_RING4_5 = 0x18A,
AMDGPU_VEGA20_DOORBELL64_UVD_RING6_7 = 0x18B,
AMDGPU_VEGA20_DOORBELL64_VCE_RING0_1 = 0x18C,
AMDGPU_VEGA20_DOORBELL64_VCE_RING2_3 = 0x18D,
AMDGPU_VEGA20_DOORBELL64_VCE_RING4_5 = 0x18E,
AMDGPU_VEGA20_DOORBELL64_VCE_RING6_7 = 0x18F,
AMDGPU_VEGA20_DOORBELL64_FIRST_NON_CP = AMDGPU_VEGA20_DOORBELL_sDMA_ENGINE0,
AMDGPU_VEGA20_DOORBELL64_LAST_NON_CP = AMDGPU_VEGA20_DOORBELL64_VCE_RING6_7,
/* kiq/kcq from second XCD. Max 8 XCDs */
AMDGPU_VEGA20_DOORBELL_XCC1_KIQ_START = 0x190,
/* 8 compute rings per GC. Max to 0x1CE */
AMDGPU_VEGA20_DOORBELL_XCC1_MEC_RING0_START = 0x197,
/* AID1 SDMA: 0x1D0 ~ 0x1F7 */
AMDGPU_VEGA20_DOORBELL_AID1_sDMA_START = 0x1D0,
AMDGPU_VEGA20_DOORBELL_MAX_ASSIGNMENT = 0x1F7,
AMDGPU_VEGA20_DOORBELL_INVALID = 0xFFFF
};
enum AMDGPU_NAVI10_DOORBELL_ASSIGNMENT {
/* Compute + GFX: 0~255 */
AMDGPU_NAVI10_DOORBELL_KIQ = 0x000,
AMDGPU_NAVI10_DOORBELL_HIQ = 0x001,
AMDGPU_NAVI10_DOORBELL_DIQ = 0x002,
AMDGPU_NAVI10_DOORBELL_MEC_RING0 = 0x003,
AMDGPU_NAVI10_DOORBELL_MEC_RING1 = 0x004,
AMDGPU_NAVI10_DOORBELL_MEC_RING2 = 0x005,
AMDGPU_NAVI10_DOORBELL_MEC_RING3 = 0x006,
AMDGPU_NAVI10_DOORBELL_MEC_RING4 = 0x007,
AMDGPU_NAVI10_DOORBELL_MEC_RING5 = 0x008,
AMDGPU_NAVI10_DOORBELL_MEC_RING6 = 0x009,
AMDGPU_NAVI10_DOORBELL_MEC_RING7 = 0x00A,
AMDGPU_NAVI10_DOORBELL_MES_RING0 = 0x00B,
AMDGPU_NAVI10_DOORBELL_MES_RING1 = 0x00C,
AMDGPU_NAVI10_DOORBELL_USERQUEUE_START = 0x00D,
AMDGPU_NAVI10_DOORBELL_USERQUEUE_END = 0x08A,
AMDGPU_NAVI10_DOORBELL_GFX_RING0 = 0x08B,
AMDGPU_NAVI10_DOORBELL_GFX_RING1 = 0x08C,
AMDGPU_NAVI10_DOORBELL_GFX_USERQUEUE_START = 0x08D,
AMDGPU_NAVI10_DOORBELL_GFX_USERQUEUE_END = 0x0FF,
/* SDMA:256~335*/
AMDGPU_NAVI10_DOORBELL_sDMA_ENGINE0 = 0x100,
AMDGPU_NAVI10_DOORBELL_sDMA_ENGINE1 = 0x10A,
AMDGPU_NAVI10_DOORBELL_sDMA_ENGINE2 = 0x114,
AMDGPU_NAVI10_DOORBELL_sDMA_ENGINE3 = 0x11E,
/* IH: 376~391 */
AMDGPU_NAVI10_DOORBELL_IH = 0x178,
/* MMSCH: 392~407
* overlap the doorbell assignment with VCN as they are mutually exclusive
* VCE engine's doorbell is 32 bit and two VCE ring share one QWORD
*/
AMDGPU_NAVI10_DOORBELL64_VCN0_1 = 0x188, /* lower 32 bits for VNC0 and upper 32 bits for VNC1 */
AMDGPU_NAVI10_DOORBELL64_VCN2_3 = 0x189,
AMDGPU_NAVI10_DOORBELL64_VCN4_5 = 0x18A,
AMDGPU_NAVI10_DOORBELL64_VCN6_7 = 0x18B,
AMDGPU_NAVI10_DOORBELL64_VCN8_9 = 0x18C,
AMDGPU_NAVI10_DOORBELL64_VCNa_b = 0x18D,
AMDGPU_NAVI10_DOORBELL64_VCNc_d = 0x18E,
AMDGPU_NAVI10_DOORBELL64_VCNe_f = 0x18F,
AMDGPU_NAVI10_DOORBELL64_VPE = 0x190,
AMDGPU_NAVI10_DOORBELL64_FIRST_NON_CP = AMDGPU_NAVI10_DOORBELL_sDMA_ENGINE0,
AMDGPU_NAVI10_DOORBELL64_LAST_NON_CP = AMDGPU_NAVI10_DOORBELL64_VPE,
AMDGPU_NAVI10_DOORBELL_MAX_ASSIGNMENT = AMDGPU_NAVI10_DOORBELL64_VPE,
AMDGPU_NAVI10_DOORBELL_INVALID = 0xFFFF
};
/*
* 64bit doorbell, offset are in QWORD, occupy 2KB doorbell space
*/
enum AMDGPU_DOORBELL64_ASSIGNMENT {
/*
* All compute related doorbells: kiq, hiq, diq, traditional compute queue, user queue, should locate in
* a continues range so that programming CP_MEC_DOORBELL_RANGE_LOWER/UPPER can cover this range.
* Compute related doorbells are allocated from 0x00 to 0x8a
*/
/* kernel scheduling */
AMDGPU_DOORBELL64_KIQ = 0x00,
/* HSA interface queue and debug queue */
AMDGPU_DOORBELL64_HIQ = 0x01,
AMDGPU_DOORBELL64_DIQ = 0x02,
/* Compute engines */
AMDGPU_DOORBELL64_MEC_RING0 = 0x03,
AMDGPU_DOORBELL64_MEC_RING1 = 0x04,
AMDGPU_DOORBELL64_MEC_RING2 = 0x05,
AMDGPU_DOORBELL64_MEC_RING3 = 0x06,
AMDGPU_DOORBELL64_MEC_RING4 = 0x07,
AMDGPU_DOORBELL64_MEC_RING5 = 0x08,
AMDGPU_DOORBELL64_MEC_RING6 = 0x09,
AMDGPU_DOORBELL64_MEC_RING7 = 0x0a,
/* User queue doorbell range (128 doorbells) */
AMDGPU_DOORBELL64_USERQUEUE_START = 0x0b,
AMDGPU_DOORBELL64_USERQUEUE_END = 0x8a,
/* Graphics engine */
AMDGPU_DOORBELL64_GFX_RING0 = 0x8b,
/*
* Other graphics doorbells can be allocated here: from 0x8c to 0xdf
* Graphics voltage island aperture 1
* default non-graphics QWORD index is 0xe0 - 0xFF inclusive
*/
/* For vega10 sriov, the sdma doorbell must be fixed as follow
* to keep the same setting with host driver, or it will
* happen conflicts
*/
AMDGPU_DOORBELL64_sDMA_ENGINE0 = 0xF0,
AMDGPU_DOORBELL64_sDMA_HI_PRI_ENGINE0 = 0xF1,
AMDGPU_DOORBELL64_sDMA_ENGINE1 = 0xF2,
AMDGPU_DOORBELL64_sDMA_HI_PRI_ENGINE1 = 0xF3,
/* Interrupt handler */
AMDGPU_DOORBELL64_IH = 0xF4, /* For legacy interrupt ring buffer */
AMDGPU_DOORBELL64_IH_RING1 = 0xF5, /* For page migration request log */
AMDGPU_DOORBELL64_IH_RING2 = 0xF6, /* For page migration translation/invalidation log */
/* VCN engine use 32 bits doorbell */
AMDGPU_DOORBELL64_VCN0_1 = 0xF8, /* lower 32 bits for VNC0 and upper 32 bits for VNC1 */
AMDGPU_DOORBELL64_VCN2_3 = 0xF9,
AMDGPU_DOORBELL64_VCN4_5 = 0xFA,
AMDGPU_DOORBELL64_VCN6_7 = 0xFB,
/* overlap the doorbell assignment with VCN as they are mutually exclusive
* VCE engine's doorbell is 32 bit and two VCE ring share one QWORD
*/
AMDGPU_DOORBELL64_UVD_RING0_1 = 0xF8,
AMDGPU_DOORBELL64_UVD_RING2_3 = 0xF9,
AMDGPU_DOORBELL64_UVD_RING4_5 = 0xFA,
AMDGPU_DOORBELL64_UVD_RING6_7 = 0xFB,
AMDGPU_DOORBELL64_VCE_RING0_1 = 0xFC,
AMDGPU_DOORBELL64_VCE_RING2_3 = 0xFD,
AMDGPU_DOORBELL64_VCE_RING4_5 = 0xFE,
AMDGPU_DOORBELL64_VCE_RING6_7 = 0xFF,
AMDGPU_DOORBELL64_FIRST_NON_CP = AMDGPU_DOORBELL64_sDMA_ENGINE0,
AMDGPU_DOORBELL64_LAST_NON_CP = AMDGPU_DOORBELL64_VCE_RING6_7,
AMDGPU_DOORBELL64_MAX_ASSIGNMENT = 0xFF,
AMDGPU_DOORBELL64_INVALID = 0xFFFF
};
enum AMDGPU_DOORBELL_ASSIGNMENT_LAYOUT1 {
/* XCC0: 0x00 ~20, XCC1: 20 ~ 2F ... */
/* KIQ/HIQ/DIQ */
AMDGPU_DOORBELL_LAYOUT1_KIQ_START = 0x000,
AMDGPU_DOORBELL_LAYOUT1_HIQ = 0x001,
AMDGPU_DOORBELL_LAYOUT1_DIQ = 0x002,
/* Compute: 0x08 ~ 0x20 */
AMDGPU_DOORBELL_LAYOUT1_MEC_RING_START = 0x008,
AMDGPU_DOORBELL_LAYOUT1_MEC_RING_END = 0x00F,
AMDGPU_DOORBELL_LAYOUT1_USERQUEUE_START = 0x010,
AMDGPU_DOORBELL_LAYOUT1_USERQUEUE_END = 0x01F,
AMDGPU_DOORBELL_LAYOUT1_XCC_RANGE = 0x020,
/* SDMA: 0x100 ~ 0x19F */
AMDGPU_DOORBELL_LAYOUT1_sDMA_ENGINE_START = 0x100,
AMDGPU_DOORBELL_LAYOUT1_sDMA_ENGINE_END = 0x19F,
/* IH: 0x1A0 ~ 0x1AF */
AMDGPU_DOORBELL_LAYOUT1_IH = 0x1A0,
/* VCN: 0x1B0 ~ 0x1E8 */
AMDGPU_DOORBELL_LAYOUT1_VCN_START = 0x1B0,
AMDGPU_DOORBELL_LAYOUT1_VCN_END = 0x1E8,
AMDGPU_DOORBELL_LAYOUT1_FIRST_NON_CP = AMDGPU_DOORBELL_LAYOUT1_sDMA_ENGINE_START,
AMDGPU_DOORBELL_LAYOUT1_LAST_NON_CP = AMDGPU_DOORBELL_LAYOUT1_VCN_END,
AMDGPU_DOORBELL_LAYOUT1_MAX_ASSIGNMENT = 0x1E8,
AMDGPU_DOORBELL_LAYOUT1_INVALID = 0xFFFF
};
#endif

79
artifacts/package_sources/tinygrad/extra/amdpci/headers/amdgpu_irq.h Normal file
View File

@@ -0,0 +1,79 @@
/*
* Copyright 2014 Advanced Micro Devices, Inc.
*
* Permission is hereby granted, free of charge, to any person obtaining a
* copy of this software and associated documentation files (the "Software"),
* to deal in the Software without restriction, including without limitation
* the rights to use, copy, modify, merge, publish, distribute, sublicense,
* and/or sell copies of the Software, and to permit persons to whom the
* Software is furnished to do so, subject to the following conditions:
*
* The above copyright notice and this permission notice shall be included in
* all copies or substantial portions of the Software.
*
* THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
* IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
* FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
* THE COPYRIGHT HOLDER(S) OR AUTHOR(S) BE LIABLE FOR ANY CLAIM, DAMAGES OR
* OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE,
* ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
* OTHER DEALINGS IN THE SOFTWARE.
*
*/
#ifndef __AMDGPU_IRQ_H__
#define __AMDGPU_IRQ_H__
// #include <linux/irqdomain.h>
// #include "soc15_ih_clientid.h"
// #include "amdgpu_ih.h"
#define bool _Bool
#define AMDGPU_MAX_IRQ_SRC_ID 0x100
#define AMDGPU_MAX_IRQ_CLIENT_ID 0x100
#define AMDGPU_IRQ_CLIENTID_LEGACY 0
#define AMDGPU_IRQ_CLIENTID_MAX SOC15_IH_CLIENTID_MAX
#define AMDGPU_IRQ_SRC_DATA_MAX_SIZE_DW 4
struct amdgpu_device;
enum amdgpu_interrupt_state {
AMDGPU_IRQ_STATE_DISABLE,
AMDGPU_IRQ_STATE_ENABLE,
};
struct amdgpu_iv_entry {
// struct amdgpu_ih_ring *ih;
unsigned client_id;
unsigned src_id;
unsigned ring_id;
unsigned vmid;
unsigned vmid_src;
uint64_t timestamp;
unsigned timestamp_src;
unsigned pasid;
unsigned node_id;
unsigned src_data[AMDGPU_IRQ_SRC_DATA_MAX_SIZE_DW];
const uint32_t *iv_entry;
};
enum interrupt_node_id_per_aid {
AID0_NODEID = 0,
XCD0_NODEID = 1,
XCD1_NODEID = 2,
AID1_NODEID = 4,
XCD2_NODEID = 5,
XCD3_NODEID = 6,
AID2_NODEID = 8,
XCD4_NODEID = 9,
XCD5_NODEID = 10,
AID3_NODEID = 12,
XCD6_NODEID = 13,
XCD7_NODEID = 14,
NODEID_MAX,
};
#endif

559
artifacts/package_sources/tinygrad/extra/amdpci/headers/amdgpu_psp.h Normal file
View File

@@ -0,0 +1,559 @@
/*
* Copyright 2016 Advanced Micro Devices, Inc.
*
* Permission is hereby granted, free of charge, to any person obtaining a
* copy of this software and associated documentation files (the "Software"),
* to deal in the Software without restriction, including without limitation
* the rights to use, copy, modify, merge, publish, distribute, sublicense,
* and/or sell copies of the Software, and to permit persons to whom the
* Software is furnished to do so, subject to the following conditions:
*
* The above copyright notice and this permission notice shall be included in
* all copies or substantial portions of the Software.
*
* THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
* IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
* FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
* THE COPYRIGHT HOLDER(S) OR AUTHOR(S) BE LIABLE FOR ANY CLAIM, DAMAGES OR
* OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE,
* ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
* OTHER DEALINGS IN THE SOFTWARE.
*
* Author: Huang Rui
*
*/
#ifndef __AMDGPU_PSP_H__
#define __AMDGPU_PSP_H__
// #include "amdgpu.h"
// #include "psp_gfx_if.h"
// #include "ta_xgmi_if.h"
// #include "ta_ras_if.h"
// #include "ta_rap_if.h"
// #include "ta_secureDisplay_if.h"
#define PSP_FENCE_BUFFER_SIZE 0x1000
#define PSP_CMD_BUFFER_SIZE 0x1000
#define PSP_1_MEG 0x100000
#define PSP_TMR_SIZE(adev) ((adev)->asic_type == CHIP_ALDEBARAN ? 0x800000 : 0x400000)
#define PSP_TMR_ALIGNMENT 0x100000
#define PSP_FW_NAME_LEN 0x24
// extern const struct attribute_group amdgpu_flash_attr_group;
enum psp_shared_mem_size {
PSP_ASD_SHARED_MEM_SIZE = 0x0,
PSP_XGMI_SHARED_MEM_SIZE = 0x4000,
PSP_RAS_SHARED_MEM_SIZE = 0x4000,
PSP_HDCP_SHARED_MEM_SIZE = 0x4000,
PSP_DTM_SHARED_MEM_SIZE = 0x4000,
PSP_RAP_SHARED_MEM_SIZE = 0x4000,
PSP_SECUREDISPLAY_SHARED_MEM_SIZE = 0x4000,
};
enum ta_type_id {
TA_TYPE_XGMI = 1,
TA_TYPE_RAS,
TA_TYPE_HDCP,
TA_TYPE_DTM,
TA_TYPE_RAP,
TA_TYPE_SECUREDISPLAY,
TA_TYPE_MAX_INDEX,
};
struct psp_context;
struct psp_xgmi_node_info;
struct psp_xgmi_topology_info;
struct psp_bin_desc;
enum psp_bootloader_cmd {
PSP_BL__LOAD_SYSDRV = 0x10000,
PSP_BL__LOAD_SOSDRV = 0x20000,
PSP_BL__LOAD_KEY_DATABASE = 0x80000,
PSP_BL__LOAD_SOCDRV = 0xB0000,
PSP_BL__LOAD_DBGDRV = 0xC0000,
PSP_BL__LOAD_HADDRV = PSP_BL__LOAD_DBGDRV,
PSP_BL__LOAD_INTFDRV = 0xD0000,
PSP_BL__LOAD_RASDRV = 0xE0000,
PSP_BL__LOAD_IPKEYMGRDRV = 0xF0000,
PSP_BL__DRAM_LONG_TRAIN = 0x100000,
PSP_BL__DRAM_SHORT_TRAIN = 0x200000,
PSP_BL__LOAD_TOS_SPL_TABLE = 0x10000000,
};
enum psp_ring_type {
PSP_RING_TYPE__INVALID = 0,
/*
* These values map to the way the PSP kernel identifies the
* rings.
*/
PSP_RING_TYPE__UM = 1, /* User mode ring (formerly called RBI) */
PSP_RING_TYPE__KM = 2 /* Kernel mode ring (formerly called GPCOM) */
};
// struct psp_ring {
// enum psp_ring_type ring_type;
// struct psp_gfx_rb_frame *ring_mem;
// uint64_t ring_mem_mc_addr;
// void *ring_mem_handle;
// uint32_t ring_size;
// uint32_t ring_wptr;
// };
/* More registers may will be supported */
enum psp_reg_prog_id {
PSP_REG_IH_RB_CNTL = 0, /* register IH_RB_CNTL */
PSP_REG_IH_RB_CNTL_RING1 = 1, /* register IH_RB_CNTL_RING1 */
PSP_REG_IH_RB_CNTL_RING2 = 2, /* register IH_RB_CNTL_RING2 */
PSP_REG_LAST
};
// struct psp_funcs {
// int (*init_microcode)(struct psp_context *psp);
// int (*wait_for_bootloader)(struct psp_context *psp);
// int (*bootloader_load_kdb)(struct psp_context *psp);
// int (*bootloader_load_spl)(struct psp_context *psp);
// int (*bootloader_load_sysdrv)(struct psp_context *psp);
// int (*bootloader_load_soc_drv)(struct psp_context *psp);
// int (*bootloader_load_intf_drv)(struct psp_context *psp);
// int (*bootloader_load_dbg_drv)(struct psp_context *psp);
// int (*bootloader_load_ras_drv)(struct psp_context *psp);
// int (*bootloader_load_ipkeymgr_drv)(struct psp_context *psp);
// int (*bootloader_load_sos)(struct psp_context *psp);
// int (*ring_create)(struct psp_context *psp,
// enum psp_ring_type ring_type);
// int (*ring_stop)(struct psp_context *psp,
// enum psp_ring_type ring_type);
// int (*ring_destroy)(struct psp_context *psp,
// enum psp_ring_type ring_type);
// bool (*smu_reload_quirk)(struct psp_context *psp);
// int (*mode1_reset)(struct psp_context *psp);
// int (*mem_training)(struct psp_context *psp, uint32_t ops);
// uint32_t (*ring_get_wptr)(struct psp_context *psp);
// void (*ring_set_wptr)(struct psp_context *psp, uint32_t value);
// int (*load_usbc_pd_fw)(struct psp_context *psp, uint64_t fw_pri_mc_addr);
// int (*read_usbc_pd_fw)(struct psp_context *psp, uint32_t *fw_ver);
// int (*update_spirom)(struct psp_context *psp, uint64_t fw_pri_mc_addr);
// int (*vbflash_stat)(struct psp_context *psp);
// int (*fatal_error_recovery_quirk)(struct psp_context *psp);
// bool (*get_ras_capability)(struct psp_context *psp);
// bool (*is_aux_sos_load_required)(struct psp_context *psp);
// };
// struct ta_funcs {
// int (*fn_ta_initialize)(struct psp_context *psp);
// int (*fn_ta_invoke)(struct psp_context *psp, uint32_t ta_cmd_id);
// int (*fn_ta_terminate)(struct psp_context *psp);
// };
#define AMDGPU_XGMI_MAX_CONNECTED_NODES 64
// struct psp_xgmi_node_info {
// uint64_t node_id;
// uint8_t num_hops;
// uint8_t is_sharing_enabled;
// enum ta_xgmi_assigned_sdma_engine sdma_engine;
// uint8_t num_links;
// struct xgmi_connected_port_num port_num[TA_XGMI__MAX_PORT_NUM];
// };
// struct psp_xgmi_topology_info {
// uint32_t num_nodes;
// struct psp_xgmi_node_info nodes[AMDGPU_XGMI_MAX_CONNECTED_NODES];
// };
// struct psp_bin_desc {
// uint32_t fw_version;
// uint32_t feature_version;
// uint32_t size_bytes;
// uint8_t *start_addr;
// };
// struct ta_mem_context {
// struct amdgpu_bo *shared_bo;
// uint64_t shared_mc_addr;
// void *shared_buf;
// enum psp_shared_mem_size shared_mem_size;
// };
// struct ta_context {
// bool initialized;
// uint32_t session_id;
// uint32_t resp_status;
// struct ta_mem_context mem_context;
// struct psp_bin_desc bin_desc;
// enum psp_gfx_cmd_id ta_load_type;
// enum ta_type_id ta_type;
// };
// struct ta_cp_context {
// struct ta_context context;
// struct mutex mutex;
// };
// struct psp_xgmi_context {
// struct ta_context context;
// struct psp_xgmi_topology_info top_info;
// bool supports_extended_data;
// uint8_t xgmi_ta_caps;
// };
// struct psp_ras_context {
// struct ta_context context;
// struct amdgpu_ras *ras;
// };
#define MEM_TRAIN_SYSTEM_SIGNATURE 0x54534942
#define GDDR6_MEM_TRAINING_DATA_SIZE_IN_BYTES 0x1000
#define GDDR6_MEM_TRAINING_OFFSET 0x8000
/*Define the VRAM size that will be encroached by BIST training.*/
#define BIST_MEM_TRAINING_ENCROACHED_SIZE 0x2000000
enum psp_memory_training_init_flag {
PSP_MEM_TRAIN_NOT_SUPPORT = 0x0,
PSP_MEM_TRAIN_SUPPORT = 0x1,
PSP_MEM_TRAIN_INIT_FAILED = 0x2,
PSP_MEM_TRAIN_RESERVE_SUCCESS = 0x4,
PSP_MEM_TRAIN_INIT_SUCCESS = 0x8,
};
enum psp_memory_training_ops {
PSP_MEM_TRAIN_SEND_LONG_MSG = 0x1,
PSP_MEM_TRAIN_SAVE = 0x2,
PSP_MEM_TRAIN_RESTORE = 0x4,
PSP_MEM_TRAIN_SEND_SHORT_MSG = 0x8,
PSP_MEM_TRAIN_COLD_BOOT = PSP_MEM_TRAIN_SEND_LONG_MSG,
PSP_MEM_TRAIN_RESUME = PSP_MEM_TRAIN_SEND_SHORT_MSG,
};
// struct psp_memory_training_context {
// /*training data size*/
// u64 train_data_size;
// /*
// * sys_cache
// * cpu virtual address
// * system memory buffer that used to store the training data.
// */
// void *sys_cache;
// /*vram offset of the p2c training data*/
// u64 p2c_train_data_offset;
// /*vram offset of the c2p training data*/
// u64 c2p_train_data_offset;
// struct amdgpu_bo *c2p_bo;
// enum psp_memory_training_init_flag init;
// u32 training_cnt;
// bool enable_mem_training;
// };
/** PSP runtime DB **/
#define PSP_RUNTIME_DB_SIZE_IN_BYTES 0x10000
#define PSP_RUNTIME_DB_OFFSET 0x100000
#define PSP_RUNTIME_DB_COOKIE_ID 0x0ed5
#define PSP_RUNTIME_DB_VER_1 0x0100
#define PSP_RUNTIME_DB_DIAG_ENTRY_MAX_COUNT 0x40
enum psp_runtime_entry_type {
PSP_RUNTIME_ENTRY_TYPE_INVALID = 0x0,
PSP_RUNTIME_ENTRY_TYPE_TEST = 0x1,
PSP_RUNTIME_ENTRY_TYPE_MGPU_COMMON = 0x2, /* Common mGPU runtime data */
PSP_RUNTIME_ENTRY_TYPE_MGPU_WAFL = 0x3, /* WAFL runtime data */
PSP_RUNTIME_ENTRY_TYPE_MGPU_XGMI = 0x4, /* XGMI runtime data */
PSP_RUNTIME_ENTRY_TYPE_BOOT_CONFIG = 0x5, /* Boot Config runtime data */
PSP_RUNTIME_ENTRY_TYPE_PPTABLE_ERR_STATUS = 0x6, /* SCPM validation data */
};
/* PSP runtime DB header */
// struct psp_runtime_data_header {
// /* determine the existence of runtime db */
// uint16_t cookie;
// /* version of runtime db */
// uint16_t version;
// };
// /* PSP runtime DB entry */
// struct psp_runtime_entry {
// /* type of runtime db entry */
// uint32_t entry_type;
// /* offset of entry in bytes */
// uint16_t offset;
// /* size of entry in bytes */
// uint16_t size;
// };
// /* PSP runtime DB directory */
// struct psp_runtime_data_directory {
// /* number of valid entries */
// uint16_t entry_count;
// /* db entries*/
// struct psp_runtime_entry entry_list[PSP_RUNTIME_DB_DIAG_ENTRY_MAX_COUNT];
// };
/* PSP runtime DB boot config feature bitmask */
enum psp_runtime_boot_cfg_feature {
BOOT_CFG_FEATURE_GECC = 0x1,
BOOT_CFG_FEATURE_TWO_STAGE_DRAM_TRAINING = 0x2,
};
/* PSP run time DB SCPM authentication defines */
enum psp_runtime_scpm_authentication {
SCPM_DISABLE = 0x0,
SCPM_ENABLE = 0x1,
SCPM_ENABLE_WITH_SCPM_ERR = 0x2,
};
/* PSP runtime DB boot config entry */
// struct psp_runtime_boot_cfg_entry {
// uint32_t boot_cfg_bitmask;
// uint32_t reserved;
// };
// /* PSP runtime DB SCPM entry */
// struct psp_runtime_scpm_entry {
// enum psp_runtime_scpm_authentication scpm_status;
// };
// struct psp_context {
// struct amdgpu_device *adev;
// struct psp_ring km_ring;
// struct psp_gfx_cmd_resp *cmd;
// const struct psp_funcs *funcs;
// const struct ta_funcs *ta_funcs;
// /* firmware buffer */
// struct amdgpu_bo *fw_pri_bo;
// uint64_t fw_pri_mc_addr;
// void *fw_pri_buf;
// /* sos firmware */
// const struct firmware *sos_fw;
// struct psp_bin_desc sys;
// struct psp_bin_desc sos;
// struct psp_bin_desc toc;
// struct psp_bin_desc kdb;
// struct psp_bin_desc spl;
// struct psp_bin_desc rl;
// struct psp_bin_desc soc_drv;
// struct psp_bin_desc intf_drv;
// struct psp_bin_desc dbg_drv;
// struct psp_bin_desc ras_drv;
// struct psp_bin_desc ipkeymgr_drv;
// /* tmr buffer */
// struct amdgpu_bo *tmr_bo;
// uint64_t tmr_mc_addr;
// /* asd firmware */
// const struct firmware *asd_fw;
// /* toc firmware */
// const struct firmware *toc_fw;
// /* cap firmware */
// const struct firmware *cap_fw;
// /* fence buffer */
// struct amdgpu_bo *fence_buf_bo;
// uint64_t fence_buf_mc_addr;
// void *fence_buf;
// /* cmd buffer */
// struct amdgpu_bo *cmd_buf_bo;
// uint64_t cmd_buf_mc_addr;
// struct psp_gfx_cmd_resp *cmd_buf_mem;
// /* fence value associated with cmd buffer */
// atomic_t fence_value;
// /* flag to mark whether gfx fw autoload is supported or not */
// bool autoload_supported;
// /* flag to mark whether psp use runtime TMR or boottime TMR */
// bool boot_time_tmr;
// /* flag to mark whether df cstate management centralized to PMFW */
// bool pmfw_centralized_cstate_management;
// /* xgmi ta firmware and buffer */
// const struct firmware *ta_fw;
// uint32_t ta_fw_version;
// uint32_t cap_fw_version;
// uint32_t cap_feature_version;
// uint32_t cap_ucode_size;
// struct ta_context asd_context;
// struct psp_xgmi_context xgmi_context;
// struct psp_ras_context ras_context;
// struct ta_cp_context hdcp_context;
// struct ta_cp_context dtm_context;
// struct ta_cp_context rap_context;
// struct ta_cp_context securedisplay_context;
// struct mutex mutex;
// struct psp_memory_training_context mem_train_ctx;
// uint32_t boot_cfg_bitmask;
// /* firmware upgrades supported */
// bool sup_pd_fw_up;
// bool sup_ifwi_up;
// char *vbflash_tmp_buf;
// size_t vbflash_image_size;
// bool vbflash_done;
// };
// struct amdgpu_psp_funcs {
// bool (*check_fw_loading_status)(struct amdgpu_device *adev,
// enum AMDGPU_UCODE_ID);
// };
// #define psp_ring_create(psp, type) (psp)->funcs->ring_create((psp), (type))
// #define psp_ring_stop(psp, type) (psp)->funcs->ring_stop((psp), (type))
// #define psp_ring_destroy(psp, type) ((psp)->funcs->ring_destroy((psp), (type)))
// #define psp_init_microcode(psp) \
// ((psp)->funcs->init_microcode ? (psp)->funcs->init_microcode((psp)) : 0)
// #define psp_bootloader_load_kdb(psp) \
// ((psp)->funcs->bootloader_load_kdb ? (psp)->funcs->bootloader_load_kdb((psp)) : 0)
// #define psp_bootloader_load_spl(psp) \
// ((psp)->funcs->bootloader_load_spl ? (psp)->funcs->bootloader_load_spl((psp)) : 0)
// #define psp_bootloader_load_sysdrv(psp) \
// ((psp)->funcs->bootloader_load_sysdrv ? (psp)->funcs->bootloader_load_sysdrv((psp)) : 0)
// #define psp_bootloader_load_soc_drv(psp) \
// ((psp)->funcs->bootloader_load_soc_drv ? (psp)->funcs->bootloader_load_soc_drv((psp)) : 0)
// #define psp_bootloader_load_intf_drv(psp) \
// ((psp)->funcs->bootloader_load_intf_drv ? (psp)->funcs->bootloader_load_intf_drv((psp)) : 0)
// #define psp_bootloader_load_dbg_drv(psp) \
// ((psp)->funcs->bootloader_load_dbg_drv ? (psp)->funcs->bootloader_load_dbg_drv((psp)) : 0)
// #define psp_bootloader_load_ras_drv(psp) \
// ((psp)->funcs->bootloader_load_ras_drv ? \
// (psp)->funcs->bootloader_load_ras_drv((psp)) : 0)
// #define psp_bootloader_load_ipkeymgr_drv(psp) \
// ((psp)->funcs->bootloader_load_ipkeymgr_drv ? \
// (psp)->funcs->bootloader_load_ipkeymgr_drv((psp)) : 0)
// #define psp_bootloader_load_sos(psp) \
// ((psp)->funcs->bootloader_load_sos ? (psp)->funcs->bootloader_load_sos((psp)) : 0)
// #define psp_smu_reload_quirk(psp) \
// ((psp)->funcs->smu_reload_quirk ? (psp)->funcs->smu_reload_quirk((psp)) : false)
// #define psp_mode1_reset(psp) \
// ((psp)->funcs->mode1_reset ? (psp)->funcs->mode1_reset((psp)) : false)
// #define psp_mem_training(psp, ops) \
// ((psp)->funcs->mem_training ? (psp)->funcs->mem_training((psp), (ops)) : 0)
// #define psp_ring_get_wptr(psp) (psp)->funcs->ring_get_wptr((psp))
// #define psp_ring_set_wptr(psp, value) (psp)->funcs->ring_set_wptr((psp), (value))
// #define psp_load_usbc_pd_fw(psp, fw_pri_mc_addr) \
// ((psp)->funcs->load_usbc_pd_fw ? \
// (psp)->funcs->load_usbc_pd_fw((psp), (fw_pri_mc_addr)) : -EINVAL)
// #define psp_read_usbc_pd_fw(psp, fw_ver) \
// ((psp)->funcs->read_usbc_pd_fw ? \
// (psp)->funcs->read_usbc_pd_fw((psp), fw_ver) : -EINVAL)
// #define psp_update_spirom(psp, fw_pri_mc_addr) \
// ((psp)->funcs->update_spirom ? \
// (psp)->funcs->update_spirom((psp), fw_pri_mc_addr) : -EINVAL)
// #define psp_vbflash_status(psp) \
// ((psp)->funcs->vbflash_stat ? \
// (psp)->funcs->vbflash_stat((psp)) : -EINVAL)
// #define psp_fatal_error_recovery_quirk(psp) \
// ((psp)->funcs->fatal_error_recovery_quirk ? \
// (psp)->funcs->fatal_error_recovery_quirk((psp)) : 0)
// #define psp_is_aux_sos_load_required(psp) \
// ((psp)->funcs->is_aux_sos_load_required ? (psp)->funcs->is_aux_sos_load_required((psp)) : 0)
// extern const struct amd_ip_funcs psp_ip_funcs;
// extern const struct amdgpu_ip_block_version psp_v3_1_ip_block;
// extern const struct amdgpu_ip_block_version psp_v10_0_ip_block;
// extern const struct amdgpu_ip_block_version psp_v11_0_ip_block;
// extern const struct amdgpu_ip_block_version psp_v11_0_8_ip_block;
// extern const struct amdgpu_ip_block_version psp_v12_0_ip_block;
// extern const struct amdgpu_ip_block_version psp_v13_0_ip_block;
// extern const struct amdgpu_ip_block_version psp_v13_0_4_ip_block;
// extern const struct amdgpu_ip_block_version psp_v14_0_ip_block;
// extern int psp_wait_for(struct psp_context *psp, uint32_t reg_index,
// uint32_t field_val, uint32_t mask, bool check_changed);
// extern int psp_wait_for_spirom_update(struct psp_context *psp, uint32_t reg_index,
// uint32_t field_val, uint32_t mask, uint32_t msec_timeout);
// int psp_execute_ip_fw_load(struct psp_context *psp,
// struct amdgpu_firmware_info *ucode);
// int psp_gpu_reset(struct amdgpu_device *adev);
// int psp_ta_init_shared_buf(struct psp_context *psp,
// struct ta_mem_context *mem_ctx);
// void psp_ta_free_shared_buf(struct ta_mem_context *mem_ctx);
// int psp_ta_unload(struct psp_context *psp, struct ta_context *context);
// int psp_ta_load(struct psp_context *psp, struct ta_context *context);
// int psp_ta_invoke(struct psp_context *psp,
// uint32_t ta_cmd_id,
// struct ta_context *context);
// int psp_xgmi_initialize(struct psp_context *psp, bool set_extended_data, bool load_ta);
// int psp_xgmi_terminate(struct psp_context *psp);
// int psp_xgmi_invoke(struct psp_context *psp, uint32_t ta_cmd_id);
// int psp_xgmi_get_hive_id(struct psp_context *psp, uint64_t *hive_id);
// int psp_xgmi_get_node_id(struct psp_context *psp, uint64_t *node_id);
// int psp_xgmi_get_topology_info(struct psp_context *psp,
// int number_devices,
// struct psp_xgmi_topology_info *topology,
// bool get_extended_data);
// int psp_xgmi_set_topology_info(struct psp_context *psp,
// int number_devices,
// struct psp_xgmi_topology_info *topology);
// int psp_ras_initialize(struct psp_context *psp);
// int psp_ras_invoke(struct psp_context *psp, uint32_t ta_cmd_id);
// int psp_ras_enable_features(struct psp_context *psp,
// union ta_ras_cmd_input *info, bool enable);
// int psp_ras_trigger_error(struct psp_context *psp,
// struct ta_ras_trigger_error_input *info, uint32_t instance_mask);
// int psp_ras_terminate(struct psp_context *psp);
// int psp_ras_query_address(struct psp_context *psp,
// struct ta_ras_query_address_input *addr_in,
// struct ta_ras_query_address_output *addr_out);
// int psp_hdcp_invoke(struct psp_context *psp, uint32_t ta_cmd_id);
// int psp_dtm_invoke(struct psp_context *psp, uint32_t ta_cmd_id);
// int psp_rap_invoke(struct psp_context *psp, uint32_t ta_cmd_id, enum ta_rap_status *status);
// int psp_securedisplay_invoke(struct psp_context *psp, uint32_t ta_cmd_id);
// int psp_rlc_autoload_start(struct psp_context *psp);
// int psp_reg_program(struct psp_context *psp, enum psp_reg_prog_id reg,
// uint32_t value);
// int psp_ring_cmd_submit(struct psp_context *psp,
// uint64_t cmd_buf_mc_addr,
// uint64_t fence_mc_addr,
// int index);
// int psp_init_asd_microcode(struct psp_context *psp,
// const char *chip_name);
// int psp_init_toc_microcode(struct psp_context *psp,
// const char *chip_name);
// int psp_init_sos_microcode(struct psp_context *psp,
// const char *chip_name);
// int psp_init_ta_microcode(struct psp_context *psp,
// const char *chip_name);
// int psp_init_cap_microcode(struct psp_context *psp,
// const char *chip_name);
// int psp_get_fw_attestation_records_addr(struct psp_context *psp,
// uint64_t *output_ptr);
// int psp_load_fw_list(struct psp_context *psp,
// struct amdgpu_firmware_info **ucode_list, int ucode_count);
// void psp_copy_fw(struct psp_context *psp, uint8_t *start_addr, uint32_t bin_size);
// int psp_spatial_partition(struct psp_context *psp, int mode);
// int is_psp_fw_valid(struct psp_bin_desc bin);
// int amdgpu_psp_wait_for_bootloader(struct amdgpu_device *adev);
// bool amdgpu_psp_get_ras_capability(struct psp_context *psp);
#endif

339
artifacts/package_sources/tinygrad/extra/amdpci/headers/amdgpu_smu.h Normal file
View File

@@ -0,0 +1,339 @@
/*
* Copyright 2019 Advanced Micro Devices, Inc.
*
* Permission is hereby granted, free of charge, to any person obtaining a
* copy of this software and associated documentation files (the "Software"),
* to deal in the Software without restriction, including without limitation
* the rights to use, copy, modify, merge, publish, distribute, sublicense,
* and/or sell copies of the Software, and to permit persons to whom the
* Software is furnished to do so, subject to the following conditions:
*
* The above copyright notice and this permission notice shall be included in
* all copies or substantial portions of the Software.
*
* THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
* IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
* FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
* THE COPYRIGHT HOLDER(S) OR AUTHOR(S) BE LIABLE FOR ANY CLAIM, DAMAGES OR
* OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE,
* ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
* OTHER DEALINGS IN THE SOFTWARE.
*/
#ifndef __AMDGPU_SMU_H__
#define __AMDGPU_SMU_H__
#define bool _Bool
#define SMU_THERMAL_MINIMUM_ALERT_TEMP 0
#define SMU_THERMAL_MAXIMUM_ALERT_TEMP 255
#define SMU_TEMPERATURE_UNITS_PER_CENTIGRADES 1000
#define SMU_FW_NAME_LEN 0x24
#define SMU_DPM_USER_PROFILE_RESTORE (1 << 0)
#define SMU_CUSTOM_FAN_SPEED_RPM (1 << 1)
#define SMU_CUSTOM_FAN_SPEED_PWM (1 << 2)
// Power Throttlers
#define SMU_THROTTLER_PPT0_BIT 0
#define SMU_THROTTLER_PPT1_BIT 1
#define SMU_THROTTLER_PPT2_BIT 2
#define SMU_THROTTLER_PPT3_BIT 3
#define SMU_THROTTLER_SPL_BIT 4
#define SMU_THROTTLER_FPPT_BIT 5
#define SMU_THROTTLER_SPPT_BIT 6
#define SMU_THROTTLER_SPPT_APU_BIT 7
// Current Throttlers
#define SMU_THROTTLER_TDC_GFX_BIT 16
#define SMU_THROTTLER_TDC_SOC_BIT 17
#define SMU_THROTTLER_TDC_MEM_BIT 18
#define SMU_THROTTLER_TDC_VDD_BIT 19
#define SMU_THROTTLER_TDC_CVIP_BIT 20
#define SMU_THROTTLER_EDC_CPU_BIT 21
#define SMU_THROTTLER_EDC_GFX_BIT 22
#define SMU_THROTTLER_APCC_BIT 23
// Temperature
#define SMU_THROTTLER_TEMP_GPU_BIT 32
#define SMU_THROTTLER_TEMP_CORE_BIT 33
#define SMU_THROTTLER_TEMP_MEM_BIT 34
#define SMU_THROTTLER_TEMP_EDGE_BIT 35
#define SMU_THROTTLER_TEMP_HOTSPOT_BIT 36
#define SMU_THROTTLER_TEMP_SOC_BIT 37
#define SMU_THROTTLER_TEMP_VR_GFX_BIT 38
#define SMU_THROTTLER_TEMP_VR_SOC_BIT 39
#define SMU_THROTTLER_TEMP_VR_MEM0_BIT 40
#define SMU_THROTTLER_TEMP_VR_MEM1_BIT 41
#define SMU_THROTTLER_TEMP_LIQUID0_BIT 42
#define SMU_THROTTLER_TEMP_LIQUID1_BIT 43
#define SMU_THROTTLER_VRHOT0_BIT 44
#define SMU_THROTTLER_VRHOT1_BIT 45
#define SMU_THROTTLER_PROCHOT_CPU_BIT 46
#define SMU_THROTTLER_PROCHOT_GFX_BIT 47
// Other
#define SMU_THROTTLER_PPM_BIT 56
#define SMU_THROTTLER_FIT_BIT 57
struct smu_hw_power_state {
unsigned int magic;
};
struct smu_power_state;
enum smu_state_ui_label {
SMU_STATE_UI_LABEL_NONE,
SMU_STATE_UI_LABEL_BATTERY,
SMU_STATE_UI_TABEL_MIDDLE_LOW,
SMU_STATE_UI_LABEL_BALLANCED,
SMU_STATE_UI_LABEL_MIDDLE_HIGHT,
SMU_STATE_UI_LABEL_PERFORMANCE,
SMU_STATE_UI_LABEL_BACO,
};
enum smu_state_classification_flag {
SMU_STATE_CLASSIFICATION_FLAG_BOOT = 0x0001,
SMU_STATE_CLASSIFICATION_FLAG_THERMAL = 0x0002,
SMU_STATE_CLASSIFICATIN_FLAG_LIMITED_POWER_SOURCE = 0x0004,
SMU_STATE_CLASSIFICATION_FLAG_RESET = 0x0008,
SMU_STATE_CLASSIFICATION_FLAG_FORCED = 0x0010,
SMU_STATE_CLASSIFICATION_FLAG_USER_3D_PERFORMANCE = 0x0020,
SMU_STATE_CLASSIFICATION_FLAG_USER_2D_PERFORMANCE = 0x0040,
SMU_STATE_CLASSIFICATION_FLAG_3D_PERFORMANCE = 0x0080,
SMU_STATE_CLASSIFICATION_FLAG_AC_OVERDIRVER_TEMPLATE = 0x0100,
SMU_STATE_CLASSIFICATION_FLAG_UVD = 0x0200,
SMU_STATE_CLASSIFICATION_FLAG_3D_PERFORMANCE_LOW = 0x0400,
SMU_STATE_CLASSIFICATION_FLAG_ACPI = 0x0800,
SMU_STATE_CLASSIFICATION_FLAG_HD2 = 0x1000,
SMU_STATE_CLASSIFICATION_FLAG_UVD_HD = 0x2000,
SMU_STATE_CLASSIFICATION_FLAG_UVD_SD = 0x4000,
SMU_STATE_CLASSIFICATION_FLAG_USER_DC_PERFORMANCE = 0x8000,
SMU_STATE_CLASSIFICATION_FLAG_DC_OVERDIRVER_TEMPLATE = 0x10000,
SMU_STATE_CLASSIFICATION_FLAG_BACO = 0x20000,
SMU_STATE_CLASSIFICATIN_FLAG_LIMITED_POWER_SOURCE2 = 0x40000,
SMU_STATE_CLASSIFICATION_FLAG_ULV = 0x80000,
SMU_STATE_CLASSIFICATION_FLAG_UVD_MVC = 0x100000,
};
struct smu_state_classification_block {
enum smu_state_ui_label ui_label;
enum smu_state_classification_flag flags;
int bios_index;
bool temporary_state;
bool to_be_deleted;
};
struct smu_state_pcie_block {
unsigned int lanes;
};
enum smu_refreshrate_source {
SMU_REFRESHRATE_SOURCE_EDID,
SMU_REFRESHRATE_SOURCE_EXPLICIT
};
struct smu_state_display_block {
bool disable_frame_modulation;
bool limit_refreshrate;
enum smu_refreshrate_source refreshrate_source;
int explicit_refreshrate;
int edid_refreshrate_index;
bool enable_vari_bright;
};
struct smu_state_memory_block {
bool dll_off;
uint8_t m3arb;
uint8_t unused[3];
};
struct smu_state_software_algorithm_block {
bool disable_load_balancing;
bool enable_sleep_for_timestamps;
};
struct smu_temperature_range {
int min;
int max;
int edge_emergency_max;
int hotspot_min;
int hotspot_crit_max;
int hotspot_emergency_max;
int mem_min;
int mem_crit_max;
int mem_emergency_max;
int software_shutdown_temp;
int software_shutdown_temp_offset;
};
struct smu_state_validation_block {
bool single_display_only;
bool disallow_on_dc;
uint8_t supported_power_levels;
};
struct smu_uvd_clocks {
uint32_t vclk;
uint32_t dclk;
};
/**
* Structure to hold a SMU Power State.
*/
enum smu_power_src_type {
SMU_POWER_SOURCE_AC,
SMU_POWER_SOURCE_DC,
SMU_POWER_SOURCE_COUNT,
};
enum smu_ppt_limit_type {
SMU_DEFAULT_PPT_LIMIT = 0,
SMU_FAST_PPT_LIMIT,
};
enum smu_ppt_limit_level {
SMU_PPT_LIMIT_MIN = -1,
SMU_PPT_LIMIT_CURRENT,
SMU_PPT_LIMIT_DEFAULT,
SMU_PPT_LIMIT_MAX,
};
enum smu_memory_pool_size {
SMU_MEMORY_POOL_SIZE_ZERO = 0,
SMU_MEMORY_POOL_SIZE_256_MB = 0x10000000,
SMU_MEMORY_POOL_SIZE_512_MB = 0x20000000,
SMU_MEMORY_POOL_SIZE_1_GB = 0x40000000,
SMU_MEMORY_POOL_SIZE_2_GB = 0x80000000,
};
enum smu_clk_type {
SMU_GFXCLK,
SMU_VCLK,
SMU_DCLK,
SMU_VCLK1,
SMU_DCLK1,
SMU_ECLK,
SMU_SOCCLK,
SMU_UCLK,
SMU_DCEFCLK,
SMU_DISPCLK,
SMU_PIXCLK,
SMU_PHYCLK,
SMU_FCLK,
SMU_SCLK,
SMU_MCLK,
SMU_PCIE,
SMU_LCLK,
SMU_OD_CCLK,
SMU_OD_SCLK,
SMU_OD_MCLK,
SMU_OD_VDDC_CURVE,
SMU_OD_RANGE,
SMU_OD_VDDGFX_OFFSET,
SMU_OD_FAN_CURVE,
SMU_OD_ACOUSTIC_LIMIT,
SMU_OD_ACOUSTIC_TARGET,
SMU_OD_FAN_TARGET_TEMPERATURE,
SMU_OD_FAN_MINIMUM_PWM,
SMU_CLK_COUNT,
};
struct smu_user_dpm_profile {
uint32_t fan_mode;
uint32_t power_limit;
uint32_t fan_speed_pwm;
uint32_t fan_speed_rpm;
uint32_t flags;
uint32_t user_od;
/* user clock state information */
uint32_t clk_mask[SMU_CLK_COUNT];
uint32_t clk_dependency;
};
#define SMU_TABLE_INIT(tables, table_id, s, a, d) \
do { \
tables[table_id].size = s; \
tables[table_id].align = a; \
tables[table_id].domain = d; \
} while (0)
struct smu_table {
uint64_t size;
uint32_t align;
uint8_t domain;
uint64_t mc_address;
void *cpu_addr;
struct amdgpu_bo *bo;
uint32_t version;
};
enum smu_perf_level_designation {
PERF_LEVEL_ACTIVITY,
PERF_LEVEL_POWER_CONTAINMENT,
};
struct smu_performance_level {
uint32_t core_clock;
uint32_t memory_clock;
uint32_t vddc;
uint32_t vddci;
uint32_t non_local_mem_freq;
uint32_t non_local_mem_width;
};
struct smu_clock_info {
uint32_t min_mem_clk;
uint32_t max_mem_clk;
uint32_t min_eng_clk;
uint32_t max_eng_clk;
uint32_t min_bus_bandwidth;
uint32_t max_bus_bandwidth;
};
struct smu_bios_boot_up_values {
uint32_t revision;
uint32_t gfxclk;
uint32_t uclk;
uint32_t socclk;
uint32_t dcefclk;
uint32_t eclk;
uint32_t vclk;
uint32_t dclk;
uint16_t vddc;
uint16_t vddci;
uint16_t mvddc;
uint16_t vdd_gfx;
uint8_t cooling_id;
uint32_t pp_table_id;
uint32_t format_revision;
uint32_t content_revision;
uint32_t fclk;
uint32_t lclk;
uint32_t firmware_caps;
};
enum smu_table_id {
SMU_TABLE_PPTABLE = 0,
SMU_TABLE_WATERMARKS,
SMU_TABLE_CUSTOM_DPM,
SMU_TABLE_DPMCLOCKS,
SMU_TABLE_AVFS,
SMU_TABLE_AVFS_PSM_DEBUG,
SMU_TABLE_AVFS_FUSE_OVERRIDE,
SMU_TABLE_PMSTATUSLOG,
SMU_TABLE_SMU_METRICS,
SMU_TABLE_DRIVER_SMU_CONFIG,
SMU_TABLE_ACTIVITY_MONITOR_COEFF,
SMU_TABLE_OVERDRIVE,
SMU_TABLE_I2C_COMMANDS,
SMU_TABLE_PACE,
SMU_TABLE_ECCINFO,
SMU_TABLE_COMBO_PPTABLE,
SMU_TABLE_WIFIBAND,
SMU_TABLE_COUNT,
};
#endif

626
artifacts/package_sources/tinygrad/extra/amdpci/headers/amdgpu_ucode.h Normal file
View File

@@ -0,0 +1,626 @@
/*
* Copyright 2012 Advanced Micro Devices, Inc.
*
* Permission is hereby granted, free of charge, to any person obtaining a
* copy of this software and associated documentation files (the "Software"),
* to deal in the Software without restriction, including without limitation
* the rights to use, copy, modify, merge, publish, distribute, sublicense,
* and/or sell copies of the Software, and to permit persons to whom the
* Software is furnished to do so, subject to the following conditions:
*
* The above copyright notice and this permission notice shall be included in
* all copies or substantial portions of the Software.
*
* THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
* IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
* FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
* THE COPYRIGHT HOLDER(S) OR AUTHOR(S) BE LIABLE FOR ANY CLAIM, DAMAGES OR
* OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE,
* ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
* OTHER DEALINGS IN THE SOFTWARE.
*
*/
#ifndef __AMDGPU_UCODE_H__
#define __AMDGPU_UCODE_H__
// #include "amdgpu_socbb.h"
#define bool _Bool
struct common_firmware_header {
uint32_t size_bytes; /* size of the entire header+image(s) in bytes */
uint32_t header_size_bytes; /* size of just the header in bytes */
uint16_t header_version_major; /* header version */
uint16_t header_version_minor; /* header version */
uint16_t ip_version_major; /* IP version */
uint16_t ip_version_minor; /* IP version */
uint32_t ucode_version;
uint32_t ucode_size_bytes; /* size of ucode in bytes */
uint32_t ucode_array_offset_bytes; /* payload offset from the start of the header */
uint32_t crc32; /* crc32 checksum of the payload */
};
/* version_major=1, version_minor=0 */
struct mc_firmware_header_v1_0 {
struct common_firmware_header header;
uint32_t io_debug_size_bytes; /* size of debug array in dwords */
uint32_t io_debug_array_offset_bytes; /* payload offset from the start of the header */
};
/* version_major=1, version_minor=0 */
struct smc_firmware_header_v1_0 {
struct common_firmware_header header;
uint32_t ucode_start_addr;
};
/* version_major=2, version_minor=0 */
struct smc_firmware_header_v2_0 {
struct smc_firmware_header_v1_0 v1_0;
uint32_t ppt_offset_bytes; /* soft pptable offset */
uint32_t ppt_size_bytes; /* soft pptable size */
};
struct smc_soft_pptable_entry {
uint32_t id;
uint32_t ppt_offset_bytes;
uint32_t ppt_size_bytes;
};
/* version_major=2, version_minor=1 */
struct smc_firmware_header_v2_1 {
struct smc_firmware_header_v1_0 v1_0;
uint32_t pptable_count;
uint32_t pptable_entry_offset;
};
struct psp_fw_legacy_bin_desc {
uint32_t fw_version;
uint32_t offset_bytes;
uint32_t size_bytes;
};
/* version_major=1, version_minor=0 */
struct psp_firmware_header_v1_0 {
struct common_firmware_header header;
struct psp_fw_legacy_bin_desc sos;
};
/* version_major=1, version_minor=1 */
struct psp_firmware_header_v1_1 {
struct psp_firmware_header_v1_0 v1_0;
struct psp_fw_legacy_bin_desc toc;
struct psp_fw_legacy_bin_desc kdb;
};
/* version_major=1, version_minor=2 */
struct psp_firmware_header_v1_2 {
struct psp_firmware_header_v1_0 v1_0;
struct psp_fw_legacy_bin_desc res;
struct psp_fw_legacy_bin_desc kdb;
};
/* version_major=1, version_minor=3 */
struct psp_firmware_header_v1_3 {
struct psp_firmware_header_v1_1 v1_1;
struct psp_fw_legacy_bin_desc spl;
struct psp_fw_legacy_bin_desc rl;
struct psp_fw_legacy_bin_desc sys_drv_aux;
struct psp_fw_legacy_bin_desc sos_aux;
};
struct psp_fw_bin_desc {
uint32_t fw_type;
uint32_t fw_version;
uint32_t offset_bytes;
uint32_t size_bytes;
};
enum psp_fw_type {
PSP_FW_TYPE_UNKOWN,
PSP_FW_TYPE_PSP_SOS,
PSP_FW_TYPE_PSP_SYS_DRV,
PSP_FW_TYPE_PSP_KDB,
PSP_FW_TYPE_PSP_TOC,
PSP_FW_TYPE_PSP_SPL,
PSP_FW_TYPE_PSP_RL,
PSP_FW_TYPE_PSP_SOC_DRV,
PSP_FW_TYPE_PSP_INTF_DRV,
PSP_FW_TYPE_PSP_DBG_DRV,
PSP_FW_TYPE_PSP_RAS_DRV,
PSP_FW_TYPE_PSP_IPKEYMGR_DRV,
PSP_FW_TYPE_MAX_INDEX,
};
/* version_major=2, version_minor=0 */
struct psp_firmware_header_v2_0 {
struct common_firmware_header header;
uint32_t psp_fw_bin_count;
struct psp_fw_bin_desc psp_fw_bin[1];
};
/* version_major=2, version_minor=1 */
struct psp_firmware_header_v2_1 {
struct common_firmware_header header;
uint32_t psp_fw_bin_count;
uint32_t psp_aux_fw_bin_index;
struct psp_fw_bin_desc psp_fw_bin[1];
};
/* version_major=1, version_minor=0 */
struct ta_firmware_header_v1_0 {
struct common_firmware_header header;
struct psp_fw_legacy_bin_desc xgmi;
struct psp_fw_legacy_bin_desc ras;
struct psp_fw_legacy_bin_desc hdcp;
struct psp_fw_legacy_bin_desc dtm;
struct psp_fw_legacy_bin_desc securedisplay;
};
enum ta_fw_type {
TA_FW_TYPE_UNKOWN,
TA_FW_TYPE_PSP_ASD,
TA_FW_TYPE_PSP_XGMI,
TA_FW_TYPE_PSP_RAS,
TA_FW_TYPE_PSP_HDCP,
TA_FW_TYPE_PSP_DTM,
TA_FW_TYPE_PSP_RAP,
TA_FW_TYPE_PSP_SECUREDISPLAY,
TA_FW_TYPE_MAX_INDEX,
};
/* version_major=2, version_minor=0 */
struct ta_firmware_header_v2_0 {
struct common_firmware_header header;
uint32_t ta_fw_bin_count;
struct psp_fw_bin_desc ta_fw_bin[1];
};
/* version_major=1, version_minor=0 */
struct gfx_firmware_header_v1_0 {
struct common_firmware_header header;
uint32_t ucode_feature_version;
uint32_t jt_offset; /* jt location */
uint32_t jt_size; /* size of jt */
};
/* version_major=2, version_minor=0 */
struct gfx_firmware_header_v2_0 {
struct common_firmware_header header;
uint32_t ucode_feature_version;
uint32_t ucode_size_bytes;
uint32_t ucode_offset_bytes;
uint32_t data_size_bytes;
uint32_t data_offset_bytes;
uint32_t ucode_start_addr_lo;
uint32_t ucode_start_addr_hi;
};
/* version_major=1, version_minor=0 */
struct mes_firmware_header_v1_0 {
struct common_firmware_header header;
uint32_t mes_ucode_version;
uint32_t mes_ucode_size_bytes;
uint32_t mes_ucode_offset_bytes;
uint32_t mes_ucode_data_version;
uint32_t mes_ucode_data_size_bytes;
uint32_t mes_ucode_data_offset_bytes;
uint32_t mes_uc_start_addr_lo;
uint32_t mes_uc_start_addr_hi;
uint32_t mes_data_start_addr_lo;
uint32_t mes_data_start_addr_hi;
};
/* version_major=1, version_minor=0 */
struct rlc_firmware_header_v1_0 {
struct common_firmware_header header;
uint32_t ucode_feature_version;
uint32_t save_and_restore_offset;
uint32_t clear_state_descriptor_offset;
uint32_t avail_scratch_ram_locations;
uint32_t master_pkt_description_offset;
};
/* version_major=2, version_minor=0 */
struct rlc_firmware_header_v2_0 {
struct common_firmware_header header;
uint32_t ucode_feature_version;
uint32_t jt_offset; /* jt location */
uint32_t jt_size; /* size of jt */
uint32_t save_and_restore_offset;
uint32_t clear_state_descriptor_offset;
uint32_t avail_scratch_ram_locations;
uint32_t reg_restore_list_size;
uint32_t reg_list_format_start;
uint32_t reg_list_format_separate_start;
uint32_t starting_offsets_start;
uint32_t reg_list_format_size_bytes; /* size of reg list format array in bytes */
uint32_t reg_list_format_array_offset_bytes; /* payload offset from the start of the header */
uint32_t reg_list_size_bytes; /* size of reg list array in bytes */
uint32_t reg_list_array_offset_bytes; /* payload offset from the start of the header */
uint32_t reg_list_format_separate_size_bytes; /* size of reg list format array in bytes */
uint32_t reg_list_format_separate_array_offset_bytes; /* payload offset from the start of the header */
uint32_t reg_list_separate_size_bytes; /* size of reg list array in bytes */
uint32_t reg_list_separate_array_offset_bytes; /* payload offset from the start of the header */
};
/* version_major=2, version_minor=1 */
struct rlc_firmware_header_v2_1 {
struct rlc_firmware_header_v2_0 v2_0;
uint32_t reg_list_format_direct_reg_list_length; /* length of direct reg list format array */
uint32_t save_restore_list_cntl_ucode_ver;
uint32_t save_restore_list_cntl_feature_ver;
uint32_t save_restore_list_cntl_size_bytes;
uint32_t save_restore_list_cntl_offset_bytes;
uint32_t save_restore_list_gpm_ucode_ver;
uint32_t save_restore_list_gpm_feature_ver;
uint32_t save_restore_list_gpm_size_bytes;
uint32_t save_restore_list_gpm_offset_bytes;
uint32_t save_restore_list_srm_ucode_ver;
uint32_t save_restore_list_srm_feature_ver;
uint32_t save_restore_list_srm_size_bytes;
uint32_t save_restore_list_srm_offset_bytes;
};
/* version_major=2, version_minor=2 */
struct rlc_firmware_header_v2_2 {
struct rlc_firmware_header_v2_1 v2_1;
uint32_t rlc_iram_ucode_size_bytes;
uint32_t rlc_iram_ucode_offset_bytes;
uint32_t rlc_dram_ucode_size_bytes;
uint32_t rlc_dram_ucode_offset_bytes;
};
/* version_major=2, version_minor=3 */
struct rlc_firmware_header_v2_3 {
struct rlc_firmware_header_v2_2 v2_2;
uint32_t rlcp_ucode_version;
uint32_t rlcp_ucode_feature_version;
uint32_t rlcp_ucode_size_bytes;
uint32_t rlcp_ucode_offset_bytes;
uint32_t rlcv_ucode_version;
uint32_t rlcv_ucode_feature_version;
uint32_t rlcv_ucode_size_bytes;
uint32_t rlcv_ucode_offset_bytes;
};
/* version_major=2, version_minor=4 */
struct rlc_firmware_header_v2_4 {
struct rlc_firmware_header_v2_3 v2_3;
uint32_t global_tap_delays_ucode_size_bytes;
uint32_t global_tap_delays_ucode_offset_bytes;
uint32_t se0_tap_delays_ucode_size_bytes;
uint32_t se0_tap_delays_ucode_offset_bytes;
uint32_t se1_tap_delays_ucode_size_bytes;
uint32_t se1_tap_delays_ucode_offset_bytes;
uint32_t se2_tap_delays_ucode_size_bytes;
uint32_t se2_tap_delays_ucode_offset_bytes;
uint32_t se3_tap_delays_ucode_size_bytes;
uint32_t se3_tap_delays_ucode_offset_bytes;
};
/* version_major=1, version_minor=0 */
struct sdma_firmware_header_v1_0 {
struct common_firmware_header header;
uint32_t ucode_feature_version;
uint32_t ucode_change_version;
uint32_t jt_offset; /* jt location */
uint32_t jt_size; /* size of jt */
};
/* version_major=1, version_minor=1 */
struct sdma_firmware_header_v1_1 {
struct sdma_firmware_header_v1_0 v1_0;
uint32_t digest_size;
};
/* version_major=2, version_minor=0 */
struct sdma_firmware_header_v2_0 {
struct common_firmware_header header;
uint32_t ucode_feature_version;
uint32_t ctx_ucode_size_bytes; /* context thread ucode size */
uint32_t ctx_jt_offset; /* context thread jt location */
uint32_t ctx_jt_size; /* context thread size of jt */
uint32_t ctl_ucode_offset;
uint32_t ctl_ucode_size_bytes; /* control thread ucode size */
uint32_t ctl_jt_offset; /* control thread jt location */
uint32_t ctl_jt_size; /* control thread size of jt */
};
/* version_major=1, version_minor=0 */
struct vpe_firmware_header_v1_0 {
struct common_firmware_header header;
uint32_t ucode_feature_version;
uint32_t ctx_ucode_size_bytes; /* context thread ucode size */
uint32_t ctx_jt_offset; /* context thread jt location */
uint32_t ctx_jt_size; /* context thread size of jt */
uint32_t ctl_ucode_offset;
uint32_t ctl_ucode_size_bytes; /* control thread ucode size */
uint32_t ctl_jt_offset; /* control thread jt location */
uint32_t ctl_jt_size; /* control thread size of jt */
};
/* version_major=1, version_minor=0 */
struct umsch_mm_firmware_header_v1_0 {
struct common_firmware_header header;
uint32_t umsch_mm_ucode_version;
uint32_t umsch_mm_ucode_size_bytes;
uint32_t umsch_mm_ucode_offset_bytes;
uint32_t umsch_mm_ucode_data_version;
uint32_t umsch_mm_ucode_data_size_bytes;
uint32_t umsch_mm_ucode_data_offset_bytes;
uint32_t umsch_mm_irq_start_addr_lo;
uint32_t umsch_mm_irq_start_addr_hi;
uint32_t umsch_mm_uc_start_addr_lo;
uint32_t umsch_mm_uc_start_addr_hi;
uint32_t umsch_mm_data_start_addr_lo;
uint32_t umsch_mm_data_start_addr_hi;
};
/* version_major=3, version_minor=0 */
struct sdma_firmware_header_v3_0 {
struct common_firmware_header header;
uint32_t ucode_feature_version;
uint32_t ucode_offset_bytes;
uint32_t ucode_size_bytes;
};
/* gpu info payload */
struct gpu_info_firmware_v1_0 {
uint32_t gc_num_se;
uint32_t gc_num_cu_per_sh;
uint32_t gc_num_sh_per_se;
uint32_t gc_num_rb_per_se;
uint32_t gc_num_tccs;
uint32_t gc_num_gprs;
uint32_t gc_num_max_gs_thds;
uint32_t gc_gs_table_depth;
uint32_t gc_gsprim_buff_depth;
uint32_t gc_parameter_cache_depth;
uint32_t gc_double_offchip_lds_buffer;
uint32_t gc_wave_size;
uint32_t gc_max_waves_per_simd;
uint32_t gc_max_scratch_slots_per_cu;
uint32_t gc_lds_size;
};
struct gpu_info_firmware_v1_1 {
struct gpu_info_firmware_v1_0 v1_0;
uint32_t num_sc_per_sh;
uint32_t num_packer_per_sc;
};
/* gpu info payload
* version_major=1, version_minor=1 */
// struct gpu_info_firmware_v1_2 {
// struct gpu_info_firmware_v1_1 v1_1;
// struct gpu_info_soc_bounding_box_v1_0 soc_bounding_box;
// };
/* version_major=1, version_minor=0 */
struct gpu_info_firmware_header_v1_0 {
struct common_firmware_header header;
uint16_t version_major; /* version */
uint16_t version_minor; /* version */
};
/* version_major=1, version_minor=0 */
struct dmcu_firmware_header_v1_0 {
struct common_firmware_header header;
uint32_t intv_offset_bytes; /* interrupt vectors offset from end of header, in bytes */
uint32_t intv_size_bytes; /* size of interrupt vectors, in bytes */
};
/* version_major=1, version_minor=0 */
struct dmcub_firmware_header_v1_0 {
struct common_firmware_header header;
uint32_t inst_const_bytes; /* size of instruction region, in bytes */
uint32_t bss_data_bytes; /* size of bss/data region, in bytes */
};
/* version_major=1, version_minor=0 */
struct imu_firmware_header_v1_0 {
struct common_firmware_header header;
uint32_t imu_iram_ucode_size_bytes;
uint32_t imu_iram_ucode_offset_bytes;
uint32_t imu_dram_ucode_size_bytes;
uint32_t imu_dram_ucode_offset_bytes;
};
/* header is fixed size */
union amdgpu_firmware_header {
struct common_firmware_header common;
struct mc_firmware_header_v1_0 mc;
struct smc_firmware_header_v1_0 smc;
struct smc_firmware_header_v2_0 smc_v2_0;
struct psp_firmware_header_v1_0 psp;
struct psp_firmware_header_v1_1 psp_v1_1;
struct psp_firmware_header_v1_3 psp_v1_3;
struct psp_firmware_header_v2_0 psp_v2_0;
struct psp_firmware_header_v2_0 psp_v2_1;
struct ta_firmware_header_v1_0 ta;
struct ta_firmware_header_v2_0 ta_v2_0;
struct gfx_firmware_header_v1_0 gfx;
struct gfx_firmware_header_v2_0 gfx_v2_0;
struct rlc_firmware_header_v1_0 rlc;
struct rlc_firmware_header_v2_0 rlc_v2_0;
struct rlc_firmware_header_v2_1 rlc_v2_1;
struct rlc_firmware_header_v2_2 rlc_v2_2;
struct rlc_firmware_header_v2_3 rlc_v2_3;
struct rlc_firmware_header_v2_4 rlc_v2_4;
struct sdma_firmware_header_v1_0 sdma;
struct sdma_firmware_header_v1_1 sdma_v1_1;
struct sdma_firmware_header_v2_0 sdma_v2_0;
struct sdma_firmware_header_v3_0 sdma_v3_0;
struct gpu_info_firmware_header_v1_0 gpu_info;
struct dmcu_firmware_header_v1_0 dmcu;
struct dmcub_firmware_header_v1_0 dmcub;
struct imu_firmware_header_v1_0 imu;
uint8_t raw[0x100];
};
#define UCODE_MAX_PSP_PACKAGING (((sizeof(union amdgpu_firmware_header) - sizeof(struct common_firmware_header) - 4) / sizeof(struct psp_fw_bin_desc)) * 2)
/*
* fw loading support
*/
enum AMDGPU_UCODE_ID {
AMDGPU_UCODE_ID_CAP = 0,
AMDGPU_UCODE_ID_SDMA0,
AMDGPU_UCODE_ID_SDMA1,
AMDGPU_UCODE_ID_SDMA2,
AMDGPU_UCODE_ID_SDMA3,
AMDGPU_UCODE_ID_SDMA4,
AMDGPU_UCODE_ID_SDMA5,
AMDGPU_UCODE_ID_SDMA6,
AMDGPU_UCODE_ID_SDMA7,
AMDGPU_UCODE_ID_SDMA_UCODE_TH0,
AMDGPU_UCODE_ID_SDMA_UCODE_TH1,
AMDGPU_UCODE_ID_SDMA_RS64,
AMDGPU_UCODE_ID_CP_CE,
AMDGPU_UCODE_ID_CP_PFP,
AMDGPU_UCODE_ID_CP_ME,
AMDGPU_UCODE_ID_CP_RS64_PFP,
AMDGPU_UCODE_ID_CP_RS64_ME,
AMDGPU_UCODE_ID_CP_RS64_MEC,
AMDGPU_UCODE_ID_CP_RS64_PFP_P0_STACK,
AMDGPU_UCODE_ID_CP_RS64_PFP_P1_STACK,
AMDGPU_UCODE_ID_CP_RS64_ME_P0_STACK,
AMDGPU_UCODE_ID_CP_RS64_ME_P1_STACK,
AMDGPU_UCODE_ID_CP_RS64_MEC_P0_STACK,
AMDGPU_UCODE_ID_CP_RS64_MEC_P1_STACK,
AMDGPU_UCODE_ID_CP_RS64_MEC_P2_STACK,
AMDGPU_UCODE_ID_CP_RS64_MEC_P3_STACK,
AMDGPU_UCODE_ID_CP_MEC1,
AMDGPU_UCODE_ID_CP_MEC1_JT,
AMDGPU_UCODE_ID_CP_MEC2,
AMDGPU_UCODE_ID_CP_MEC2_JT,
AMDGPU_UCODE_ID_CP_MES,
AMDGPU_UCODE_ID_CP_MES_DATA,
AMDGPU_UCODE_ID_CP_MES1,
AMDGPU_UCODE_ID_CP_MES1_DATA,
AMDGPU_UCODE_ID_IMU_I,
AMDGPU_UCODE_ID_IMU_D,
AMDGPU_UCODE_ID_GLOBAL_TAP_DELAYS,
AMDGPU_UCODE_ID_SE0_TAP_DELAYS,
AMDGPU_UCODE_ID_SE1_TAP_DELAYS,
AMDGPU_UCODE_ID_SE2_TAP_DELAYS,
AMDGPU_UCODE_ID_SE3_TAP_DELAYS,
AMDGPU_UCODE_ID_RLC_RESTORE_LIST_CNTL,
AMDGPU_UCODE_ID_RLC_RESTORE_LIST_GPM_MEM,
AMDGPU_UCODE_ID_RLC_RESTORE_LIST_SRM_MEM,
AMDGPU_UCODE_ID_RLC_IRAM,
AMDGPU_UCODE_ID_RLC_DRAM,
AMDGPU_UCODE_ID_RLC_P,
AMDGPU_UCODE_ID_RLC_V,
AMDGPU_UCODE_ID_RLC_G,
AMDGPU_UCODE_ID_STORAGE,
AMDGPU_UCODE_ID_SMC,
AMDGPU_UCODE_ID_PPTABLE,
AMDGPU_UCODE_ID_UVD,
AMDGPU_UCODE_ID_UVD1,
AMDGPU_UCODE_ID_VCE,
AMDGPU_UCODE_ID_VCN,
AMDGPU_UCODE_ID_VCN1,
AMDGPU_UCODE_ID_DMCU_ERAM,
AMDGPU_UCODE_ID_DMCU_INTV,
AMDGPU_UCODE_ID_VCN0_RAM,
AMDGPU_UCODE_ID_VCN1_RAM,
AMDGPU_UCODE_ID_DMCUB,
AMDGPU_UCODE_ID_VPE_CTX,
AMDGPU_UCODE_ID_VPE_CTL,
AMDGPU_UCODE_ID_VPE,
AMDGPU_UCODE_ID_UMSCH_MM_UCODE,
AMDGPU_UCODE_ID_UMSCH_MM_DATA,
AMDGPU_UCODE_ID_UMSCH_MM_CMD_BUFFER,
AMDGPU_UCODE_ID_P2S_TABLE,
AMDGPU_UCODE_ID_JPEG_RAM,
AMDGPU_UCODE_ID_ISP,
AMDGPU_UCODE_ID_MAXIMUM,
};
/* engine firmware status */
enum AMDGPU_UCODE_STATUS {
AMDGPU_UCODE_STATUS_INVALID,
AMDGPU_UCODE_STATUS_NOT_LOADED,
AMDGPU_UCODE_STATUS_LOADED,
};
enum amdgpu_firmware_load_type {
AMDGPU_FW_LOAD_DIRECT = 0,
AMDGPU_FW_LOAD_PSP,
AMDGPU_FW_LOAD_SMU,
AMDGPU_FW_LOAD_RLC_BACKDOOR_AUTO,
};
/* conform to smu_ucode_xfer_cz.h */
#define AMDGPU_SDMA0_UCODE_LOADED 0x00000001
#define AMDGPU_SDMA1_UCODE_LOADED 0x00000002
#define AMDGPU_CPCE_UCODE_LOADED 0x00000004
#define AMDGPU_CPPFP_UCODE_LOADED 0x00000008
#define AMDGPU_CPME_UCODE_LOADED 0x00000010
#define AMDGPU_CPMEC1_UCODE_LOADED 0x00000020
#define AMDGPU_CPMEC2_UCODE_LOADED 0x00000040
#define AMDGPU_CPRLC_UCODE_LOADED 0x00000100
/* amdgpu firmware info */
struct amdgpu_firmware_info {
/* ucode ID */
enum AMDGPU_UCODE_ID ucode_id;
/* request_firmware */
const struct firmware *fw;
/* starting mc address */
uint64_t mc_addr;
/* kernel linear address */
void *kaddr;
/* ucode_size_bytes */
uint32_t ucode_size;
/* starting tmr mc address */
uint32_t tmr_mc_addr_lo;
uint32_t tmr_mc_addr_hi;
};
// struct amdgpu_firmware {
// struct amdgpu_firmware_info ucode[AMDGPU_UCODE_ID_MAXIMUM];
// enum amdgpu_firmware_load_type load_type;
// struct amdgpu_bo *fw_buf;
// unsigned int fw_size;
// unsigned int max_ucodes;
// /* firmwares are loaded by psp instead of smu from vega10 */
// const struct amdgpu_psp_funcs *funcs;
// struct amdgpu_bo *rbuf;
// struct mutex mutex;
// /* gpu info firmware data pointer */
// const struct firmware *gpu_info_fw;
// void *fw_buf_ptr;
// uint64_t fw_buf_mc;
// };
// void amdgpu_ucode_print_mc_hdr(const struct common_firmware_header *hdr);
// void amdgpu_ucode_print_smc_hdr(const struct common_firmware_header *hdr);
// void amdgpu_ucode_print_imu_hdr(const struct common_firmware_header *hdr);
// void amdgpu_ucode_print_gfx_hdr(const struct common_firmware_header *hdr);
// void amdgpu_ucode_print_rlc_hdr(const struct common_firmware_header *hdr);
// void amdgpu_ucode_print_sdma_hdr(const struct common_firmware_header *hdr);
// void amdgpu_ucode_print_psp_hdr(const struct common_firmware_header *hdr);
// void amdgpu_ucode_print_gpu_info_hdr(const struct common_firmware_header *hdr);
// int amdgpu_ucode_request(struct amdgpu_device *adev, const struct firmware **fw,
// const char *fw_name);
// void amdgpu_ucode_release(const struct firmware **fw);
// bool amdgpu_ucode_hdr_version(union amdgpu_firmware_header *hdr,
// uint16_t hdr_major, uint16_t hdr_minor);
// int amdgpu_ucode_init_bo(struct amdgpu_device *adev);
// int amdgpu_ucode_create_bo(struct amdgpu_device *adev);
// int amdgpu_ucode_sysfs_init(struct amdgpu_device *adev);
// void amdgpu_ucode_free_bo(struct amdgpu_device *adev);
// void amdgpu_ucode_sysfs_fini(struct amdgpu_device *adev);
// enum amdgpu_firmware_load_type
// amdgpu_ucode_get_load_type(struct amdgpu_device *adev, int load_type);
// const char *amdgpu_ucode_name(enum AMDGPU_UCODE_ID ucode_id);
// void amdgpu_ucode_ip_version_decode(struct amdgpu_device *adev, int block_type, char *ucode_prefix, int len);
#endif

665
artifacts/package_sources/tinygrad/extra/amdpci/headers/amdgpu_vm.h Normal file
View File

@@ -0,0 +1,665 @@
/*
* Copyright 2016 Advanced Micro Devices, Inc.
*
* Permission is hereby granted, free of charge, to any person obtaining a
* copy of this software and associated documentation files (the "Software"),
* to deal in the Software without restriction, including without limitation
* the rights to use, copy, modify, merge, publish, distribute, sublicense,
* and/or sell copies of the Software, and to permit persons to whom the
* Software is furnished to do so, subject to the following conditions:
*
* The above copyright notice and this permission notice shall be included in
* all copies or substantial portions of the Software.
*
* THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
* IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
* FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
* THE COPYRIGHT HOLDER(S) OR AUTHOR(S) BE LIABLE FOR ANY CLAIM, DAMAGES OR
* OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE,
* ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
* OTHER DEALINGS IN THE SOFTWARE.
*
* Authors: Christian König
*/
#ifndef __AMDGPU_VM_H__
#define __AMDGPU_VM_H__
// #include <linux/idr.h>
// #include <linux/kfifo.h>
// #include <linux/rbtree.h>
// #include <drm/gpu_scheduler.h>
// #include <drm/drm_file.h>
// #include <drm/ttm/ttm_bo.h>
// #include <linux/sched/mm.h>
// #include "amdgpu_sync.h"
// #include "amdgpu_ring.h"
// #include "amdgpu_ids.h"
// struct drm_exec;
// struct amdgpu_bo_va;
// struct amdgpu_job;
// struct amdgpu_bo_list_entry;
// struct amdgpu_bo_vm;
// struct amdgpu_mem_stats;
/*
* GPUVM handling
*/
/* Maximum number of PTEs the hardware can write with one command */
#define AMDGPU_VM_MAX_UPDATE_SIZE 0x3FFFF
/* number of entries in page table */
#define AMDGPU_VM_PTE_COUNT(adev) (1 << (adev)->vm_manager.block_size)
#define AMDGPU_PTE_VALID (1ULL << 0)
#define AMDGPU_PTE_SYSTEM (1ULL << 1)
#define AMDGPU_PTE_SNOOPED (1ULL << 2)
/* RV+ */
#define AMDGPU_PTE_TMZ (1ULL << 3)
/* VI only */
#define AMDGPU_PTE_EXECUTABLE (1ULL << 4)
#define AMDGPU_PTE_READABLE (1ULL << 5)
#define AMDGPU_PTE_WRITEABLE (1ULL << 6)
#define AMDGPU_PTE_FRAG(x) ((x & 0x1fULL) << 7)
/* TILED for VEGA10, reserved for older ASICs */
#define AMDGPU_PTE_PRT (1ULL << 51)
/* PDE is handled as PTE for VEGA10 */
#define AMDGPU_PDE_PTE (1ULL << 54)
#define AMDGPU_PTE_LOG (1ULL << 55)
/* PTE is handled as PDE for VEGA10 (Translate Further) */
#define AMDGPU_PTE_TF (1ULL << 56)
/* MALL noalloc for sienna_cichlid, reserved for older ASICs */
#define AMDGPU_PTE_NOALLOC (1ULL << 58)
/* PDE Block Fragment Size for VEGA10 */
#define AMDGPU_PDE_BFS(a) ((uint64_t)a << 59)
/* Flag combination to set no-retry with TF disabled */
#define AMDGPU_VM_NORETRY_FLAGS (AMDGPU_PTE_EXECUTABLE | AMDGPU_PDE_PTE | \
AMDGPU_PTE_TF)
/* Flag combination to set no-retry with TF enabled */
#define AMDGPU_VM_NORETRY_FLAGS_TF (AMDGPU_PTE_VALID | AMDGPU_PTE_SYSTEM | \
AMDGPU_PTE_PRT)
/* For GFX9 */
#define AMDGPU_PTE_MTYPE_VG10_SHIFT(mtype) ((uint64_t)(mtype) << 57)
#define AMDGPU_PTE_MTYPE_VG10_MASK AMDGPU_PTE_MTYPE_VG10_SHIFT(3ULL)
#define AMDGPU_PTE_MTYPE_VG10(flags, mtype) \
(((uint64_t)(flags) & (~AMDGPU_PTE_MTYPE_VG10_MASK)) | \
AMDGPU_PTE_MTYPE_VG10_SHIFT(mtype))
#define AMDGPU_MTYPE_NC 0
#define AMDGPU_MTYPE_CC 2
#define AMDGPU_PTE_DEFAULT_ATC (AMDGPU_PTE_SYSTEM \
| AMDGPU_PTE_SNOOPED \
| AMDGPU_PTE_EXECUTABLE \
| AMDGPU_PTE_READABLE \
| AMDGPU_PTE_WRITEABLE \
| AMDGPU_PTE_MTYPE_VG10(AMDGPU_MTYPE_CC))
/* gfx10 */
#define AMDGPU_PTE_MTYPE_NV10_SHIFT(mtype) ((uint64_t)(mtype) << 48)
#define AMDGPU_PTE_MTYPE_NV10_MASK AMDGPU_PTE_MTYPE_NV10_SHIFT(7ULL)
#define AMDGPU_PTE_MTYPE_NV10(flags, mtype) \
(((uint64_t)(flags) & (~AMDGPU_PTE_MTYPE_NV10_MASK)) | \
AMDGPU_PTE_MTYPE_NV10_SHIFT(mtype))
/* gfx12 */
#define AMDGPU_PTE_PRT_GFX12 (1ULL << 56)
#define AMDGPU_PTE_PRT_FLAG(adev) \
((amdgpu_ip_version((adev), GC_HWIP, 0) >= IP_VERSION(12, 0, 0)) ? AMDGPU_PTE_PRT_GFX12 : AMDGPU_PTE_PRT)
#define AMDGPU_PTE_MTYPE_GFX12_SHIFT(mtype) ((uint64_t)(mtype) << 54)
#define AMDGPU_PTE_MTYPE_GFX12_MASK AMDGPU_PTE_MTYPE_GFX12_SHIFT(3ULL)
#define AMDGPU_PTE_MTYPE_GFX12(flags, mtype) \
(((uint64_t)(flags) & (~AMDGPU_PTE_MTYPE_GFX12_MASK)) | \
AMDGPU_PTE_MTYPE_GFX12_SHIFT(mtype))
#define AMDGPU_PTE_IS_PTE (1ULL << 63)
/* PDE Block Fragment Size for gfx v12 */
#define AMDGPU_PDE_BFS_GFX12(a) ((uint64_t)((a) & 0x1fULL) << 58)
#define AMDGPU_PDE_BFS_FLAG(adev, a) \
((amdgpu_ip_version((adev), GC_HWIP, 0) >= IP_VERSION(12, 0, 0)) ? AMDGPU_PDE_BFS_GFX12(a) : AMDGPU_PDE_BFS(a))
/* PDE is handled as PTE for gfx v12 */
#define AMDGPU_PDE_PTE_GFX12 (1ULL << 63)
#define AMDGPU_PDE_PTE_FLAG(adev) \
((amdgpu_ip_version((adev), GC_HWIP, 0) >= IP_VERSION(12, 0, 0)) ? AMDGPU_PDE_PTE_GFX12 : AMDGPU_PDE_PTE)
/* How to program VM fault handling */
#define AMDGPU_VM_FAULT_STOP_NEVER 0
#define AMDGPU_VM_FAULT_STOP_FIRST 1
#define AMDGPU_VM_FAULT_STOP_ALWAYS 2
/* How much VRAM be reserved for page tables */
#define AMDGPU_VM_RESERVED_VRAM (8ULL << 20)
/*
* max number of VMHUB
* layout: max 8 GFXHUB + 4 MMHUB0 + 1 MMHUB1
*/
#define AMDGPU_MAX_VMHUBS 13
#define AMDGPU_GFXHUB_START 0
#define AMDGPU_MMHUB0_START 8
#define AMDGPU_MMHUB1_START 12
#define AMDGPU_GFXHUB(x) (AMDGPU_GFXHUB_START + (x))
#define AMDGPU_MMHUB0(x) (AMDGPU_MMHUB0_START + (x))
#define AMDGPU_MMHUB1(x) (AMDGPU_MMHUB1_START + (x))
#define AMDGPU_IS_GFXHUB(x) ((x) >= AMDGPU_GFXHUB_START && (x) < AMDGPU_MMHUB0_START)
#define AMDGPU_IS_MMHUB0(x) ((x) >= AMDGPU_MMHUB0_START && (x) < AMDGPU_MMHUB1_START)
#define AMDGPU_IS_MMHUB1(x) ((x) >= AMDGPU_MMHUB1_START && (x) < AMDGPU_MAX_VMHUBS)
/* Reserve space at top/bottom of address space for kernel use */
#define AMDGPU_VA_RESERVED_CSA_SIZE (2ULL << 20)
#define AMDGPU_VA_RESERVED_CSA_START(adev) (((adev)->vm_manager.max_pfn \
<< AMDGPU_GPU_PAGE_SHIFT) \
- AMDGPU_VA_RESERVED_CSA_SIZE)
#define AMDGPU_VA_RESERVED_SEQ64_SIZE (2ULL << 20)
#define AMDGPU_VA_RESERVED_SEQ64_START(adev) (AMDGPU_VA_RESERVED_CSA_START(adev) \
- AMDGPU_VA_RESERVED_SEQ64_SIZE)
#define AMDGPU_VA_RESERVED_TRAP_SIZE (2ULL << 12)
#define AMDGPU_VA_RESERVED_TRAP_START(adev) (AMDGPU_VA_RESERVED_SEQ64_START(adev) \
- AMDGPU_VA_RESERVED_TRAP_SIZE)
#define AMDGPU_VA_RESERVED_BOTTOM (1ULL << 16)
#define AMDGPU_VA_RESERVED_TOP (AMDGPU_VA_RESERVED_TRAP_SIZE + \
AMDGPU_VA_RESERVED_SEQ64_SIZE + \
AMDGPU_VA_RESERVED_CSA_SIZE)
/* See vm_update_mode */
#define AMDGPU_VM_USE_CPU_FOR_GFX (1 << 0)
#define AMDGPU_VM_USE_CPU_FOR_COMPUTE (1 << 1)
/* VMPT level enumerate, and the hiberachy is:
* PDB2->PDB1->PDB0->PTB
*/
enum amdgpu_vm_level {
AMDGPU_VM_PDB2,
AMDGPU_VM_PDB1,
AMDGPU_VM_PDB0,
AMDGPU_VM_PTB
};
// /* base structure for tracking BO usage in a VM */
// struct amdgpu_vm_bo_base {
// /* constant after initialization */
// struct amdgpu_vm *vm;
// struct amdgpu_bo *bo;
// /* protected by bo being reserved */
// struct amdgpu_vm_bo_base *next;
// /* protected by spinlock */
// struct list_head vm_status;
// /* protected by the BO being reserved */
// bool moved;
// };
// /* provided by hw blocks that can write ptes, e.g., sdma */
// struct amdgpu_vm_pte_funcs {
// /* number of dw to reserve per operation */
// unsigned copy_pte_num_dw;
// /* copy pte entries from GART */
// void (*copy_pte)(struct amdgpu_ib *ib,
// uint64_t pe, uint64_t src,
// unsigned count);
// /* write pte one entry at a time with addr mapping */
// void (*write_pte)(struct amdgpu_ib *ib, uint64_t pe,
// uint64_t value, unsigned count,
// uint32_t incr);
// /* for linear pte/pde updates without addr mapping */
// void (*set_pte_pde)(struct amdgpu_ib *ib,
// uint64_t pe,
// uint64_t addr, unsigned count,
// uint32_t incr, uint64_t flags);
// };
// struct amdgpu_task_info {
// char process_name[TASK_COMM_LEN];
// char task_name[TASK_COMM_LEN];
// pid_t pid;
// pid_t tgid;
// struct kref refcount;
// };
// /**
// * struct amdgpu_vm_update_params
// *
// * Encapsulate some VM table update parameters to reduce
// * the number of function parameters
// *
// */
// struct amdgpu_vm_update_params {
// /**
// * @adev: amdgpu device we do this update for
// */
// struct amdgpu_device *adev;
// /**
// * @vm: optional amdgpu_vm we do this update for
// */
// struct amdgpu_vm *vm;
// /**
// * @immediate: if changes should be made immediately
// */
// bool immediate;
// /**
// * @unlocked: true if the root BO is not locked
// */
// bool unlocked;
// /**
// * @pages_addr:
// *
// * DMA addresses to use for mapping
// */
// dma_addr_t *pages_addr;
// /**
// * @job: job to used for hw submission
// */
// struct amdgpu_job *job;
// /**
// * @num_dw_left: number of dw left for the IB
// */
// unsigned int num_dw_left;
// /**
// * @needs_flush: true whenever we need to invalidate the TLB
// */
// bool needs_flush;
// /**
// * @allow_override: true for memory that is not uncached: allows MTYPE
// * to be overridden for NUMA local memory.
// */
// bool allow_override;
// /**
// * @tlb_flush_waitlist: temporary storage for BOs until tlb_flush
// */
// struct list_head tlb_flush_waitlist;
// };
// struct amdgpu_vm_update_funcs {
// int (*map_table)(struct amdgpu_bo_vm *bo);
// int (*prepare)(struct amdgpu_vm_update_params *p, struct dma_resv *resv,
// enum amdgpu_sync_mode sync_mode);
// int (*update)(struct amdgpu_vm_update_params *p,
// struct amdgpu_bo_vm *bo, uint64_t pe, uint64_t addr,
// unsigned count, uint32_t incr, uint64_t flags);
// int (*commit)(struct amdgpu_vm_update_params *p,
// struct dma_fence **fence);
// };
// struct amdgpu_vm_fault_info {
// /* fault address */
// uint64_t addr;
// /* fault status register */
// uint32_t status;
// /* which vmhub? gfxhub, mmhub, etc. */
// unsigned int vmhub;
// };
// struct amdgpu_vm {
// /* tree of virtual addresses mapped */
// #ifndef HAVE_TREE_INSERT_HAVE_RB_ROOT_CACHED
// struct rb_root va;
// #else
// struct rb_root_cached va;
// #endif
// /* Lock to prevent eviction while we are updating page tables
// * use vm_eviction_lock/unlock(vm)
// */
// struct mutex eviction_lock;
// bool evicting;
// unsigned int saved_flags;
// /* Lock to protect vm_bo add/del/move on all lists of vm */
// spinlock_t status_lock;
// /* Per-VM and PT BOs who needs a validation */
// struct list_head evicted;
// /* BOs for user mode queues that need a validation */
// struct list_head evicted_user;
// /* PT BOs which relocated and their parent need an update */
// struct list_head relocated;
// /* per VM BOs moved, but not yet updated in the PT */
// struct list_head moved;
// /* All BOs of this VM not currently in the state machine */
// struct list_head idle;
// /* regular invalidated BOs, but not yet updated in the PT */
// struct list_head invalidated;
// /* BO mappings freed, but not yet updated in the PT */
// struct list_head freed;
// /* BOs which are invalidated, has been updated in the PTs */
// struct list_head done;
// /* PT BOs scheduled to free and fill with zero if vm_resv is not hold */
// struct list_head pt_freed;
// struct work_struct pt_free_work;
// /* contains the page directory */
// struct amdgpu_vm_bo_base root;
// struct dma_fence *last_update;
// /* Scheduler entities for page table updates */
// struct drm_sched_entity immediate;
// struct drm_sched_entity delayed;
// /* Last finished delayed update */
// atomic64_t tlb_seq;
// struct dma_fence *last_tlb_flush;
// atomic64_t kfd_last_flushed_seq;
// uint64_t tlb_fence_context;
// /* How many times we had to re-generate the page tables */
// uint64_t generation;
// /* Last unlocked submission to the scheduler entities */
// struct dma_fence *last_unlocked;
// unsigned int pasid;
// bool reserved_vmid[AMDGPU_MAX_VMHUBS];
// /* Flag to indicate if VM tables are updated by CPU or GPU (SDMA) */
// bool use_cpu_for_update;
// /* Functions to use for VM table updates */
// const struct amdgpu_vm_update_funcs *update_funcs;
// /* Up to 128 pending retry page faults */
// DECLARE_KFIFO(faults, u64, 128);
// /* Points to the KFD process VM info */
// struct amdkfd_process_info *process_info;
// /* List node in amdkfd_process_info.vm_list_head */
// struct list_head vm_list_node;
// /* Valid while the PD is reserved or fenced */
// uint64_t pd_phys_addr;
// /* Some basic info about the task */
// struct amdgpu_task_info *task_info;
// /* Store positions of group of BOs */
// struct ttm_lru_bulk_move lru_bulk_move;
// /* Flag to indicate if VM is used for compute */
// bool is_compute_context;
// /* Memory partition number, -1 means any partition */
// int8_t mem_id;
// /* cached fault info */
// struct amdgpu_vm_fault_info fault_info;
// };
// struct amdgpu_vm_manager {
// /* Handling of VMIDs */
// struct amdgpu_vmid_mgr id_mgr[AMDGPU_MAX_VMHUBS];
// unsigned int first_kfd_vmid;
// bool concurrent_flush;
// /* Handling of VM fences */
// u64 fence_context;
// unsigned seqno[AMDGPU_MAX_RINGS];
// uint64_t max_pfn;
// uint32_t num_level;
// uint32_t block_size;
// uint32_t fragment_size;
// enum amdgpu_vm_level root_level;
// /* vram base address for page table entry */
// u64 vram_base_offset;
// /* vm pte handling */
// const struct amdgpu_vm_pte_funcs *vm_pte_funcs;
// struct drm_gpu_scheduler *vm_pte_scheds[AMDGPU_MAX_RINGS];
// unsigned vm_pte_num_scheds;
// struct amdgpu_ring *page_fault;
// /* partial resident texture handling */
// spinlock_t prt_lock;
// atomic_t num_prt_users;
// /* controls how VM page tables are updated for Graphics and Compute.
// * BIT0[= 0] Graphics updated by SDMA [= 1] by CPU
// * BIT1[= 0] Compute updated by SDMA [= 1] by CPU
// */
// int vm_update_mode;
// /* PASID to VM mapping, will be used in interrupt context to
// * look up VM of a page fault
// */
// #ifdef HAVE_STRUCT_XARRAY
// struct xarray pasids;
// #else
// struct idr pasid_idr;
// spinlock_t pasid_lock;
// #endif
// /* Global registration of recent page fault information */
// struct amdgpu_vm_fault_info fault_info;
// };
// struct amdgpu_bo_va_mapping;
// #define amdgpu_vm_copy_pte(adev, ib, pe, src, count) ((adev)->vm_manager.vm_pte_funcs->copy_pte((ib), (pe), (src), (count)))
// #define amdgpu_vm_write_pte(adev, ib, pe, value, count, incr) ((adev)->vm_manager.vm_pte_funcs->write_pte((ib), (pe), (value), (count), (incr)))
// #define amdgpu_vm_set_pte_pde(adev, ib, pe, addr, count, incr, flags) ((adev)->vm_manager.vm_pte_funcs->set_pte_pde((ib), (pe), (addr), (count), (incr), (flags)))
// extern const struct amdgpu_vm_update_funcs amdgpu_vm_cpu_funcs;
// extern const struct amdgpu_vm_update_funcs amdgpu_vm_sdma_funcs;
// void amdgpu_vm_manager_init(struct amdgpu_device *adev);
// void amdgpu_vm_manager_fini(struct amdgpu_device *adev);
// int amdgpu_vm_set_pasid(struct amdgpu_device *adev, struct amdgpu_vm *vm,
// u32 pasid);
// long amdgpu_vm_wait_idle(struct amdgpu_vm *vm, long timeout);
// int amdgpu_vm_init(struct amdgpu_device *adev, struct amdgpu_vm *vm, int32_t xcp_id);
// int amdgpu_vm_make_compute(struct amdgpu_device *adev, struct amdgpu_vm *vm);
// void amdgpu_vm_release_compute(struct amdgpu_device *adev, struct amdgpu_vm *vm);
// void amdgpu_vm_fini(struct amdgpu_device *adev, struct amdgpu_vm *vm);
// int amdgpu_vm_lock_pd(struct amdgpu_vm *vm, struct drm_exec *exec,
// unsigned int num_fences);
// bool amdgpu_vm_ready(struct amdgpu_vm *vm);
// uint64_t amdgpu_vm_generation(struct amdgpu_device *adev, struct amdgpu_vm *vm);
// int amdgpu_vm_validate(struct amdgpu_device *adev, struct amdgpu_vm *vm,
// struct ww_acquire_ctx *ticket,
// int (*callback)(void *p, struct amdgpu_bo *bo),
// void *param);
// int amdgpu_vm_flush(struct amdgpu_ring *ring, struct amdgpu_job *job, bool need_pipe_sync);
// int amdgpu_vm_update_pdes(struct amdgpu_device *adev,
// struct amdgpu_vm *vm, bool immediate);
// int amdgpu_vm_clear_freed(struct amdgpu_device *adev,
// struct amdgpu_vm *vm,
// struct dma_fence **fence);
// int amdgpu_vm_handle_moved(struct amdgpu_device *adev,
// struct amdgpu_vm *vm,
// struct ww_acquire_ctx *ticket);
// int amdgpu_vm_flush_compute_tlb(struct amdgpu_device *adev,
// struct amdgpu_vm *vm,
// uint32_t flush_type,
// uint32_t xcc_mask);
// void amdgpu_vm_bo_base_init(struct amdgpu_vm_bo_base *base,
// struct amdgpu_vm *vm, struct amdgpu_bo *bo);
// int amdgpu_vm_update_range(struct amdgpu_device *adev, struct amdgpu_vm *vm,
// bool immediate, bool unlocked, bool flush_tlb, bool allow_override,
// struct dma_resv *resv, uint64_t start, uint64_t last,
// uint64_t flags, uint64_t offset, uint64_t vram_base,
// struct ttm_resource *res, dma_addr_t *pages_addr,
// struct dma_fence **fence);
// int amdgpu_vm_bo_update(struct amdgpu_device *adev,
// struct amdgpu_bo_va *bo_va,
// bool clear);
// bool amdgpu_vm_evictable(struct amdgpu_bo *bo);
// void amdgpu_vm_bo_invalidate(struct amdgpu_device *adev,
// struct amdgpu_bo *bo, bool evicted);
// uint64_t amdgpu_vm_map_gart(const dma_addr_t *pages_addr, uint64_t addr);
// struct amdgpu_bo_va *amdgpu_vm_bo_find(struct amdgpu_vm *vm,
// struct amdgpu_bo *bo);
// struct amdgpu_bo_va *amdgpu_vm_bo_add(struct amdgpu_device *adev,
// struct amdgpu_vm *vm,
// struct amdgpu_bo *bo);
// int amdgpu_vm_bo_map(struct amdgpu_device *adev,
// struct amdgpu_bo_va *bo_va,
// uint64_t addr, uint64_t offset,
// uint64_t size, uint64_t flags);
// int amdgpu_vm_bo_replace_map(struct amdgpu_device *adev,
// struct amdgpu_bo_va *bo_va,
// uint64_t addr, uint64_t offset,
// uint64_t size, uint64_t flags);
// int amdgpu_vm_bo_unmap(struct amdgpu_device *adev,
// struct amdgpu_bo_va *bo_va,
// uint64_t addr);
// int amdgpu_vm_bo_clear_mappings(struct amdgpu_device *adev,
// struct amdgpu_vm *vm,
// uint64_t saddr, uint64_t size);
// struct amdgpu_bo_va_mapping *amdgpu_vm_bo_lookup_mapping(struct amdgpu_vm *vm,
// uint64_t addr);
// void amdgpu_vm_bo_trace_cs(struct amdgpu_vm *vm, struct ww_acquire_ctx *ticket);
// void amdgpu_vm_bo_del(struct amdgpu_device *adev,
// struct amdgpu_bo_va *bo_va);
// void amdgpu_vm_adjust_size(struct amdgpu_device *adev, uint32_t min_vm_size,
// uint32_t fragment_size_default, unsigned max_level,
// unsigned max_bits);
// int amdgpu_vm_ioctl(struct drm_device *dev, void *data, struct drm_file *filp);
// bool amdgpu_vm_need_pipeline_sync(struct amdgpu_ring *ring,
// struct amdgpu_job *job);
// void amdgpu_vm_check_compute_bug(struct amdgpu_device *adev);
// struct amdgpu_task_info *
// amdgpu_vm_get_task_info_pasid(struct amdgpu_device *adev, u32 pasid);
// struct amdgpu_task_info *
// amdgpu_vm_get_task_info_vm(struct amdgpu_vm *vm);
// void amdgpu_vm_put_task_info(struct amdgpu_task_info *task_info);
// bool amdgpu_vm_handle_fault(struct amdgpu_device *adev, u32 pasid,
// u32 vmid, u32 node_id, uint64_t addr,
// bool write_fault);
// void amdgpu_vm_set_task_info(struct amdgpu_vm *vm);
// void amdgpu_vm_move_to_lru_tail(struct amdgpu_device *adev,
// struct amdgpu_vm *vm);
// void amdgpu_vm_get_memory(struct amdgpu_vm *vm,
// struct amdgpu_mem_stats *stats);
// int amdgpu_vm_pt_clear(struct amdgpu_device *adev, struct amdgpu_vm *vm,
// struct amdgpu_bo_vm *vmbo, bool immediate);
// int amdgpu_vm_pt_create(struct amdgpu_device *adev, struct amdgpu_vm *vm,
// int level, bool immediate, struct amdgpu_bo_vm **vmbo,
// int32_t xcp_id);
// void amdgpu_vm_pt_free_root(struct amdgpu_device *adev, struct amdgpu_vm *vm);
// int amdgpu_vm_pde_update(struct amdgpu_vm_update_params *params,
// struct amdgpu_vm_bo_base *entry);
// int amdgpu_vm_ptes_update(struct amdgpu_vm_update_params *params,
// uint64_t start, uint64_t end,
// uint64_t dst, uint64_t flags);
// void amdgpu_vm_pt_free_work(struct work_struct *work);
// void amdgpu_vm_pt_free_list(struct amdgpu_device *adev,
// struct amdgpu_vm_update_params *params);
// #if defined(CONFIG_DEBUG_FS)
// void amdgpu_debugfs_vm_bo_info(struct amdgpu_vm *vm, struct seq_file *m);
// #endif
// int amdgpu_vm_pt_map_tables(struct amdgpu_device *adev, struct amdgpu_vm *vm);
// bool amdgpu_vm_is_bo_always_valid(struct amdgpu_vm *vm, struct amdgpu_bo *bo);
// /**
// * amdgpu_vm_tlb_seq - return tlb flush sequence number
// * @vm: the amdgpu_vm structure to query
// *
// * Returns the tlb flush sequence number which indicates that the VM TLBs needs
// * to be invalidated whenever the sequence number change.
// */
// static inline uint64_t amdgpu_vm_tlb_seq(struct amdgpu_vm *vm)
// {
// unsigned long flags;
// spinlock_t *lock;
// /*
// * Workaround to stop racing between the fence signaling and handling
// * the cb. The lock is static after initially setting it up, just make
// * sure that the dma_fence structure isn't freed up.
// */
// rcu_read_lock();
// lock = vm->last_tlb_flush->lock;
// rcu_read_unlock();
// spin_lock_irqsave(lock, flags);
// spin_unlock_irqrestore(lock, flags);
// return atomic64_read(&vm->tlb_seq);
// }
// /*
// * vm eviction_lock can be taken in MMU notifiers. Make sure no reclaim-FS
// * happens while holding this lock anywhere to prevent deadlocks when
// * an MMU notifier runs in reclaim-FS context.
// */
// static inline void amdgpu_vm_eviction_lock(struct amdgpu_vm *vm)
// {
// mutex_lock(&vm->eviction_lock);
// vm->saved_flags = memalloc_noreclaim_save();
// }
// static inline bool amdgpu_vm_eviction_trylock(struct amdgpu_vm *vm)
// {
// if (mutex_trylock(&vm->eviction_lock)) {
// vm->saved_flags = memalloc_noreclaim_save();
// return true;
// }
// return false;
// }
// static inline void amdgpu_vm_eviction_unlock(struct amdgpu_vm *vm)
// {
// memalloc_noreclaim_restore(vm->saved_flags);
// mutex_unlock(&vm->eviction_lock);
// }
// void amdgpu_vm_update_fault_cache(struct amdgpu_device *adev,
// unsigned int pasid,
// uint64_t addr,
// uint32_t status,
// unsigned int vmhub);
// void amdgpu_vm_tlb_fence_create(struct amdgpu_device *adev,
// struct amdgpu_vm *vm,
// struct dma_fence **fence);
#endif

600
artifacts/package_sources/tinygrad/extra/amdpci/headers/discovery.h Normal file
View File

@@ -0,0 +1,600 @@
/*
* Copyright 2018 Advanced Micro Devices, Inc.
*
* Permission is hereby granted, free of charge, to any person obtaining a
* copy of this software and associated documentation files (the "Software"),
* to deal in the Software without restriction, including without limitation
* the rights to use, copy, modify, merge, publish, distribute, sublicense,
* and/or sell copies of the Software, and to permit persons to whom the
* Software is furnished to do so, subject to the following conditions:
*
* The above copyright notice and this permission notice shall be included in
* all copies or substantial portions of the Software.
*
* THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
* IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
* FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
* THE COPYRIGHT HOLDER(S) OR AUTHOR(S) BE LIABLE FOR ANY CLAIM, DAMAGES OR
* OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE,
* ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
* OTHER DEALINGS IN THE SOFTWARE.
*
*/
#ifndef _DISCOVERY_H_
#define _DISCOVERY_H_
#define PSP_HEADER_SIZE 256
#define BINARY_SIGNATURE 0x28211407
#define DISCOVERY_TABLE_SIGNATURE 0x53445049
#define GC_TABLE_ID 0x4347
#define HARVEST_TABLE_SIGNATURE 0x56524148
#define VCN_INFO_TABLE_ID 0x004E4356
#define MALL_INFO_TABLE_ID 0x4C4C414D
#define NPS_INFO_TABLE_ID 0x0053504E
typedef enum {
IP_DISCOVERY = 0,
GC,
HARVEST_INFO,
VCN_INFO,
MALL_INFO,
NPS_INFO,
TOTAL_TABLES = 6
} table;
#pragma pack(1)
typedef struct table_info
{
uint16_t offset; /* Byte offset */
uint16_t checksum; /* Byte sum of the table */
uint16_t size; /* Table size */
uint16_t padding;
} table_info;
typedef struct binary_header
{
/* psp structure should go at the top of this structure */
uint32_t binary_signature; /* 0x7, 0x14, 0x21, 0x28 */
uint16_t version_major;
uint16_t version_minor;
uint16_t binary_checksum; /* Byte sum of the binary after this field */
uint16_t binary_size; /* Binary Size*/
table_info table_list[TOTAL_TABLES];
} binary_header;
typedef struct die_info
{
uint16_t die_id;
uint16_t die_offset; /* Points to the corresponding die_header structure */
} die_info;
typedef struct ip_discovery_header
{
uint32_t signature; /* Table Signature */
uint16_t version; /* Table Version */
uint16_t size; /* Table Size */
uint32_t id; /* Table ID */
uint16_t num_dies; /* Number of Dies */
die_info die_info[16]; /* list die information for up to 16 dies */
union {
uint16_t padding[1]; /* version <= 3 */
struct { /* version == 4 */
uint8_t base_addr_64_bit : 1; /* ip structures are using 64 bit base address */
uint8_t reserved : 7;
uint8_t reserved2;
};
};
} ip_discovery_header;
typedef struct ip
{
uint16_t hw_id; /* Hardware ID */
uint8_t number_instance; /* instance of the IP */
uint8_t num_base_address; /* Number of Base Addresses */
uint8_t major; /* HCID Major */
uint8_t minor; /* HCID Minor */
uint8_t revision; /* HCID Revision */
#if defined(__BIG_ENDIAN)
uint8_t reserved : 4; /* Placeholder field */
uint8_t harvest : 4; /* Harvest */
#else
uint8_t harvest : 4; /* Harvest */
uint8_t reserved : 4; /* Placeholder field */
#endif
uint32_t base_address[]; /* variable number of Addresses */
} ip;
typedef struct ip_v3
{
uint16_t hw_id; /* Hardware ID */
uint8_t instance_number; /* Instance number for the IP */
uint8_t num_base_address; /* Number of base addresses*/
uint8_t major; /* Hardware ID.major version */
uint8_t minor; /* Hardware ID.minor version */
uint8_t revision; /* Hardware ID.revision version */
#if defined(__BIG_ENDIAN)
uint8_t variant : 4; /* HW variant */
uint8_t sub_revision : 4; /* HCID Sub-Revision */
#else
uint8_t sub_revision : 4; /* HCID Sub-Revision */
uint8_t variant : 4; /* HW variant */
#endif
uint32_t base_address[]; /* Base Address list. Corresponds to the num_base_address field*/
} ip_v3;
typedef struct ip_v4 {
uint16_t hw_id; /* Hardware ID */
uint8_t instance_number; /* Instance number for the IP */
uint8_t num_base_address; /* Number of base addresses*/
uint8_t major; /* Hardware ID.major version */
uint8_t minor; /* Hardware ID.minor version */
uint8_t revision; /* Hardware ID.revision version */
#if defined(LITTLEENDIAN_CPU)
uint8_t sub_revision : 4; /* HCID Sub-Revision */
uint8_t variant : 4; /* HW variant */
#elif defined(BIGENDIAN_CPU)
uint8_t variant : 4; /* HW variant */
uint8_t sub_revision : 4; /* HCID Sub-Revision */
#endif
} ip_v4;
typedef struct die_header
{
uint16_t die_id;
uint16_t num_ips;
} die_header;
typedef struct ip_structure
{
ip_discovery_header* header;
struct die
{
die_header *die_header;
union
{
ip *ip_list;
ip_v3 *ip_v3_list;
ip_v4 *ip_v4_list;
}; /* IP list. Variable size*/
} die;
} ip_structure;
struct gpu_info_header {
uint32_t table_id; /* table ID */
uint16_t version_major; /* table version */
uint16_t version_minor; /* table version */
uint32_t size; /* size of the entire header+data in bytes */
};
struct gc_info_v1_0 {
struct gpu_info_header header;
uint32_t gc_num_se;
uint32_t gc_num_wgp0_per_sa;
uint32_t gc_num_wgp1_per_sa;
uint32_t gc_num_rb_per_se;
uint32_t gc_num_gl2c;
uint32_t gc_num_gprs;
uint32_t gc_num_max_gs_thds;
uint32_t gc_gs_table_depth;
uint32_t gc_gsprim_buff_depth;
uint32_t gc_parameter_cache_depth;
uint32_t gc_double_offchip_lds_buffer;
uint32_t gc_wave_size;
uint32_t gc_max_waves_per_simd;
uint32_t gc_max_scratch_slots_per_cu;
uint32_t gc_lds_size;
uint32_t gc_num_sc_per_se;
uint32_t gc_num_sa_per_se;
uint32_t gc_num_packer_per_sc;
uint32_t gc_num_gl2a;
};
struct gc_info_v1_1 {
struct gpu_info_header header;
uint32_t gc_num_se;
uint32_t gc_num_wgp0_per_sa;
uint32_t gc_num_wgp1_per_sa;
uint32_t gc_num_rb_per_se;
uint32_t gc_num_gl2c;
uint32_t gc_num_gprs;
uint32_t gc_num_max_gs_thds;
uint32_t gc_gs_table_depth;
uint32_t gc_gsprim_buff_depth;
uint32_t gc_parameter_cache_depth;
uint32_t gc_double_offchip_lds_buffer;
uint32_t gc_wave_size;
uint32_t gc_max_waves_per_simd;
uint32_t gc_max_scratch_slots_per_cu;
uint32_t gc_lds_size;
uint32_t gc_num_sc_per_se;
uint32_t gc_num_sa_per_se;
uint32_t gc_num_packer_per_sc;
uint32_t gc_num_gl2a;
uint32_t gc_num_tcp_per_sa;
uint32_t gc_num_sdp_interface;
uint32_t gc_num_tcps;
};
struct gc_info_v1_2 {
struct gpu_info_header header;
uint32_t gc_num_se;
uint32_t gc_num_wgp0_per_sa;
uint32_t gc_num_wgp1_per_sa;
uint32_t gc_num_rb_per_se;
uint32_t gc_num_gl2c;
uint32_t gc_num_gprs;
uint32_t gc_num_max_gs_thds;
uint32_t gc_gs_table_depth;
uint32_t gc_gsprim_buff_depth;
uint32_t gc_parameter_cache_depth;
uint32_t gc_double_offchip_lds_buffer;
uint32_t gc_wave_size;
uint32_t gc_max_waves_per_simd;
uint32_t gc_max_scratch_slots_per_cu;
uint32_t gc_lds_size;
uint32_t gc_num_sc_per_se;
uint32_t gc_num_sa_per_se;
uint32_t gc_num_packer_per_sc;
uint32_t gc_num_gl2a;
uint32_t gc_num_tcp_per_sa;
uint32_t gc_num_sdp_interface;
uint32_t gc_num_tcps;
uint32_t gc_num_tcp_per_wpg;
uint32_t gc_tcp_l1_size;
uint32_t gc_num_sqc_per_wgp;
uint32_t gc_l1_instruction_cache_size_per_sqc;
uint32_t gc_l1_data_cache_size_per_sqc;
uint32_t gc_gl1c_per_sa;
uint32_t gc_gl1c_size_per_instance;
uint32_t gc_gl2c_per_gpu;
};
struct gc_info_v1_3 {
struct gpu_info_header header;
uint32_t gc_num_se;
uint32_t gc_num_wgp0_per_sa;
uint32_t gc_num_wgp1_per_sa;
uint32_t gc_num_rb_per_se;
uint32_t gc_num_gl2c;
uint32_t gc_num_gprs;
uint32_t gc_num_max_gs_thds;
uint32_t gc_gs_table_depth;
uint32_t gc_gsprim_buff_depth;
uint32_t gc_parameter_cache_depth;
uint32_t gc_double_offchip_lds_buffer;
uint32_t gc_wave_size;
uint32_t gc_max_waves_per_simd;
uint32_t gc_max_scratch_slots_per_cu;
uint32_t gc_lds_size;
uint32_t gc_num_sc_per_se;
uint32_t gc_num_sa_per_se;
uint32_t gc_num_packer_per_sc;
uint32_t gc_num_gl2a;
uint32_t gc_num_tcp_per_sa;
uint32_t gc_num_sdp_interface;
uint32_t gc_num_tcps;
uint32_t gc_num_tcp_per_wpg;
uint32_t gc_tcp_l1_size;
uint32_t gc_num_sqc_per_wgp;
uint32_t gc_l1_instruction_cache_size_per_sqc;
uint32_t gc_l1_data_cache_size_per_sqc;
uint32_t gc_gl1c_per_sa;
uint32_t gc_gl1c_size_per_instance;
uint32_t gc_gl2c_per_gpu;
uint32_t gc_tcp_size_per_cu;
uint32_t gc_tcp_cache_line_size;
uint32_t gc_instruction_cache_size_per_sqc;
uint32_t gc_instruction_cache_line_size;
uint32_t gc_scalar_data_cache_size_per_sqc;
uint32_t gc_scalar_data_cache_line_size;
uint32_t gc_tcc_size;
uint32_t gc_tcc_cache_line_size;
};
struct gc_info_v2_0 {
struct gpu_info_header header;
uint32_t gc_num_se;
uint32_t gc_num_cu_per_sh;
uint32_t gc_num_sh_per_se;
uint32_t gc_num_rb_per_se;
uint32_t gc_num_tccs;
uint32_t gc_num_gprs;
uint32_t gc_num_max_gs_thds;
uint32_t gc_gs_table_depth;
uint32_t gc_gsprim_buff_depth;
uint32_t gc_parameter_cache_depth;
uint32_t gc_double_offchip_lds_buffer;
uint32_t gc_wave_size;
uint32_t gc_max_waves_per_simd;
uint32_t gc_max_scratch_slots_per_cu;
uint32_t gc_lds_size;
uint32_t gc_num_sc_per_se;
uint32_t gc_num_packer_per_sc;
};
struct gc_info_v2_1 {
struct gpu_info_header header;
uint32_t gc_num_se;
uint32_t gc_num_cu_per_sh;
uint32_t gc_num_sh_per_se;
uint32_t gc_num_rb_per_se;
uint32_t gc_num_tccs;
uint32_t gc_num_gprs;
uint32_t gc_num_max_gs_thds;
uint32_t gc_gs_table_depth;
uint32_t gc_gsprim_buff_depth;
uint32_t gc_parameter_cache_depth;
uint32_t gc_double_offchip_lds_buffer;
uint32_t gc_wave_size;
uint32_t gc_max_waves_per_simd;
uint32_t gc_max_scratch_slots_per_cu;
uint32_t gc_lds_size;
uint32_t gc_num_sc_per_se;
uint32_t gc_num_packer_per_sc;
/* new for v2_1 */
uint32_t gc_num_tcp_per_sh;
uint32_t gc_tcp_size_per_cu;
uint32_t gc_num_sdp_interface;
uint32_t gc_num_cu_per_sqc;
uint32_t gc_instruction_cache_size_per_sqc;
uint32_t gc_scalar_data_cache_size_per_sqc;
uint32_t gc_tcc_size;
};
typedef struct harvest_info_header {
uint32_t signature; /* Table Signature */
uint32_t version; /* Table Version */
} harvest_info_header;
typedef struct harvest_info {
uint16_t hw_id; /* Hardware ID */
uint8_t number_instance; /* Instance of the IP */
uint8_t reserved; /* Reserved for alignment */
} harvest_info;
typedef struct harvest_table {
harvest_info_header header;
harvest_info list[32];
} harvest_table;
struct mall_info_header {
uint32_t table_id; /* table ID */
uint16_t version_major; /* table version */
uint16_t version_minor; /* table version */
uint32_t size_bytes; /* size of the entire header+data in bytes */
};
struct mall_info_v1_0 {
struct mall_info_header header;
uint32_t mall_size_per_m;
uint32_t m_s_present;
uint32_t m_half_use;
uint32_t m_mall_config;
uint32_t reserved[5];
};
struct mall_info_v2_0 {
struct mall_info_header header;
uint32_t mall_size_per_umc;
uint32_t reserved[8];
};
#define VCN_INFO_TABLE_MAX_NUM_INSTANCES 4
struct vcn_info_header {
uint32_t table_id; /* table ID */
uint16_t version_major; /* table version */
uint16_t version_minor; /* table version */
uint32_t size_bytes; /* size of the entire header+data in bytes */
};
struct vcn_instance_info_v1_0
{
uint32_t instance_num; /* VCN IP instance number. 0 - VCN0; 1 - VCN1 etc*/
union _fuse_data {
struct {
uint32_t av1_disabled : 1;
uint32_t vp9_disabled : 1;
uint32_t hevc_disabled : 1;
uint32_t h264_disabled : 1;
uint32_t reserved : 28;
} bits;
uint32_t all_bits;
} fuse_data;
uint32_t reserved[2];
};
struct vcn_info_v1_0 {
struct vcn_info_header header;
uint32_t num_of_instances; /* number of entries used in instance_info below*/
struct vcn_instance_info_v1_0 instance_info[VCN_INFO_TABLE_MAX_NUM_INSTANCES];
uint32_t reserved[4];
};
#define NPS_INFO_TABLE_MAX_NUM_INSTANCES 12
struct nps_info_header {
uint32_t table_id; /* table ID */
uint16_t version_major; /* table version */
uint16_t version_minor; /* table version */
uint32_t size_bytes; /* size of the entire header+data in bytes = 0x000000D4 (212) */
};
struct nps_instance_info_v1_0 {
uint64_t base_address;
uint64_t limit_address;
};
struct nps_info_v1_0 {
struct nps_info_header header;
uint32_t nps_type;
uint32_t count;
struct nps_instance_info_v1_0
instance_info[NPS_INFO_TABLE_MAX_NUM_INSTANCES];
};
enum amd_hw_ip_block_type {
GC_HWIP = 1,
HDP_HWIP,
SDMA0_HWIP,
SDMA1_HWIP,
SDMA2_HWIP,
SDMA3_HWIP,
SDMA4_HWIP,
SDMA5_HWIP,
SDMA6_HWIP,
SDMA7_HWIP,
LSDMA_HWIP,
MMHUB_HWIP,
ATHUB_HWIP,
NBIO_HWIP,
MP0_HWIP,
MP1_HWIP,
UVD_HWIP,
VCN_HWIP = UVD_HWIP,
JPEG_HWIP = VCN_HWIP,
VCN1_HWIP,
VCE_HWIP,
VPE_HWIP,
DF_HWIP,
DCE_HWIP,
OSSSYS_HWIP,
SMUIO_HWIP,
PWR_HWIP,
NBIF_HWIP,
THM_HWIP,
CLK_HWIP,
UMC_HWIP,
RSMU_HWIP,
XGMI_HWIP,
DCI_HWIP,
PCIE_HWIP,
ISP_HWIP,
MAX_HWIP
};
#define HWIP_MAX_INSTANCE 44
#define HW_ID_MAX 300
// HW ID
#define MP1_HWID 1
#define MP2_HWID 2
#define THM_HWID 3
#define SMUIO_HWID 4
#define FUSE_HWID 5
#define CLKA_HWID 6
#define PWR_HWID 10
#define GC_HWID 11
#define UVD_HWID 12
#define VCN_HWID UVD_HWID
#define AUDIO_AZ_HWID 13
#define ACP_HWID 14
#define DCI_HWID 15
#define DMU_HWID 271
#define DCO_HWID 16
#define DIO_HWID 272
#define XDMA_HWID 17
#define DCEAZ_HWID 18
#define DAZ_HWID 274
#define SDPMUX_HWID 19
#define NTB_HWID 20
#define VPE_HWID 21
#define IOHC_HWID 24
#define L2IMU_HWID 28
#define VCE_HWID 32
#define MMHUB_HWID 34
#define ATHUB_HWID 35
#define DBGU_NBIO_HWID 36
#define DFX_HWID 37
#define DBGU0_HWID 38
#define DBGU1_HWID 39
#define OSSSYS_HWID 40
#define HDP_HWID 41
#define SDMA0_HWID 42
#define SDMA1_HWID 43
#define ISP_HWID 44
#define DBGU_IO_HWID 45
#define DF_HWID 46
#define CLKB_HWID 47
#define FCH_HWID 48
#define DFX_DAP_HWID 49
#define L1IMU_PCIE_HWID 50
#define L1IMU_NBIF_HWID 51
#define L1IMU_IOAGR_HWID 52
#define L1IMU3_HWID 53
#define L1IMU4_HWID 54
#define L1IMU5_HWID 55
#define L1IMU6_HWID 56
#define L1IMU7_HWID 57
#define L1IMU8_HWID 58
#define L1IMU9_HWID 59
#define L1IMU10_HWID 60
#define L1IMU11_HWID 61
#define L1IMU12_HWID 62
#define L1IMU13_HWID 63
#define L1IMU14_HWID 64
#define L1IMU15_HWID 65
#define WAFLC_HWID 66
#define FCH_USB_PD_HWID 67
#define SDMA2_HWID 68
#define SDMA3_HWID 69
#define PCIE_HWID 70
#define PCS_HWID 80
#define DDCL_HWID 89
#define SST_HWID 90
#define LSDMA_HWID 91
#define IOAGR_HWID 100
#define NBIF_HWID 108
#define IOAPIC_HWID 124
#define SYSTEMHUB_HWID 128
#define NTBCCP_HWID 144
#define UMC_HWID 150
#define SATA_HWID 168
#define USB_HWID 170
#define CCXSEC_HWID 176
#define XGMI_HWID 200
#define XGBE_HWID 216
#define MP0_HWID 255
static int hw_id_map[MAX_HWIP] = {
[GC_HWIP] = GC_HWID,
[HDP_HWIP] = HDP_HWID,
[SDMA0_HWIP] = SDMA0_HWID,
[SDMA1_HWIP] = SDMA1_HWID,
[SDMA2_HWIP] = SDMA2_HWID,
[SDMA3_HWIP] = SDMA3_HWID,
[LSDMA_HWIP] = LSDMA_HWID,
[MMHUB_HWIP] = MMHUB_HWID,
[ATHUB_HWIP] = ATHUB_HWID,
[NBIO_HWIP] = NBIF_HWID,
[MP0_HWIP] = MP0_HWID,
[MP1_HWIP] = MP1_HWID,
[UVD_HWIP] = UVD_HWID,
[VCE_HWIP] = VCE_HWID,
[DF_HWIP] = DF_HWID,
[DCE_HWIP] = DMU_HWID,
[OSSSYS_HWIP] = OSSSYS_HWID,
[SMUIO_HWIP] = SMUIO_HWID,
[PWR_HWIP] = PWR_HWID,
[NBIF_HWIP] = NBIF_HWID,
[THM_HWIP] = THM_HWID,
[CLK_HWIP] = CLKA_HWID,
[UMC_HWIP] = UMC_HWID,
[XGMI_HWIP] = XGMI_HWID,
[DCI_HWIP] = DCI_HWID,
[PCIE_HWIP] = PCIE_HWID,
[VPE_HWIP] = VPE_HWID,
[ISP_HWIP] = ISP_HWID,
};
#endif

471
artifacts/package_sources/tinygrad/extra/amdpci/headers/psp_gfx_if.h Normal file
View File

@@ -0,0 +1,471 @@
/*
* Copyright 2017 Advanced Micro Devices, Inc.
*
* Permission is hereby granted, free of charge, to any person obtaining a
* copy of this software and associated documentation files (the "Software"),
* to deal in the Software without restriction, including without limitation
* the rights to use, copy, modify, merge, publish, distribute, sublicense,
* and/or sell copies of the Software, and to permit persons to whom the
* Software is furnished to do so, subject to the following conditions:
*
* The above copyright notice and this permission notice shall be included in
* all copies or substantial portions of the Software.
*
* THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
* IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
* FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
* THE COPYRIGHT HOLDER(S) OR AUTHOR(S) BE LIABLE FOR ANY CLAIM, DAMAGES OR
* OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE,
* ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
* OTHER DEALINGS IN THE SOFTWARE.
*
*/
#ifndef _PSP_TEE_GFX_IF_H_
#define _PSP_TEE_GFX_IF_H_
#define PSP_GFX_CMD_BUF_VERSION 0x00000001
#define GFX_CMD_STATUS_MASK 0x0000FFFF
#define GFX_CMD_ID_MASK 0x000F0000
#define GFX_CMD_RESERVED_MASK 0x7FF00000
#define GFX_CMD_RESPONSE_MASK 0x80000000
/* USBC PD FW version retrieval command */
#define C2PMSG_CMD_GFX_USB_PD_FW_VER 0x2000000
/* TEE Gfx Command IDs for the register interface.
* Command ID must be between 0x00010000 and 0x000F0000.
*/
enum psp_gfx_crtl_cmd_id
{
GFX_CTRL_CMD_ID_INIT_RBI_RING = 0x00010000, /* initialize RBI ring */
GFX_CTRL_CMD_ID_INIT_GPCOM_RING = 0x00020000, /* initialize GPCOM ring */
GFX_CTRL_CMD_ID_DESTROY_RINGS = 0x00030000, /* destroy rings */
GFX_CTRL_CMD_ID_CAN_INIT_RINGS = 0x00040000, /* is it allowed to initialized the rings */
GFX_CTRL_CMD_ID_ENABLE_INT = 0x00050000, /* enable PSP-to-Gfx interrupt */
GFX_CTRL_CMD_ID_DISABLE_INT = 0x00060000, /* disable PSP-to-Gfx interrupt */
GFX_CTRL_CMD_ID_MODE1_RST = 0x00070000, /* trigger the Mode 1 reset */
GFX_CTRL_CMD_ID_GBR_IH_SET = 0x00080000, /* set Gbr IH_RB_CNTL registers */
GFX_CTRL_CMD_ID_CONSUME_CMD = 0x00090000, /* send interrupt to psp for updating write pointer of vf */
GFX_CTRL_CMD_ID_DESTROY_GPCOM_RING = 0x000C0000, /* destroy GPCOM ring */
GFX_CTRL_CMD_ID_MAX = 0x000F0000, /* max command ID */
};
/*-----------------------------------------------------------------------------
NOTE: All physical addresses used in this interface are actually
GPU Virtual Addresses.
*/
/* Control registers of the TEE Gfx interface. These are located in
* SRBM-to-PSP mailbox registers (total 8 registers).
*/
struct psp_gfx_ctrl
{
volatile unsigned int cmd_resp; /* +0 Command/Response register for Gfx commands */
volatile unsigned int rbi_wptr; /* +4 Write pointer (index) of RBI ring */
volatile unsigned int rbi_rptr; /* +8 Read pointer (index) of RBI ring */
volatile unsigned int gpcom_wptr; /* +12 Write pointer (index) of GPCOM ring */
volatile unsigned int gpcom_rptr; /* +16 Read pointer (index) of GPCOM ring */
volatile unsigned int ring_addr_lo; /* +20 bits [31:0] of GPU Virtual of ring buffer (VMID=0)*/
volatile unsigned int ring_addr_hi; /* +24 bits [63:32] of GPU Virtual of ring buffer (VMID=0) */
volatile unsigned int ring_buf_size; /* +28 Ring buffer size (in bytes) */
};
/* Response flag is set in the command when command is completed by PSP.
* Used in the GFX_CTRL.CmdResp.
* When PSP GFX I/F is initialized, the flag is set.
*/
#define GFX_FLAG_RESPONSE 0x80000000
/* TEE Gfx Command IDs for the ring buffer interface. */
enum psp_gfx_cmd_id
{
GFX_CMD_ID_LOAD_TA = 0x00000001, /* load TA */
GFX_CMD_ID_UNLOAD_TA = 0x00000002, /* unload TA */
GFX_CMD_ID_INVOKE_CMD = 0x00000003, /* send command to TA */
GFX_CMD_ID_LOAD_ASD = 0x00000004, /* load ASD Driver */
GFX_CMD_ID_SETUP_TMR = 0x00000005, /* setup TMR region */
GFX_CMD_ID_LOAD_IP_FW = 0x00000006, /* load HW IP FW */
GFX_CMD_ID_DESTROY_TMR = 0x00000007, /* destroy TMR region */
GFX_CMD_ID_SAVE_RESTORE = 0x00000008, /* save/restore HW IP FW */
GFX_CMD_ID_SETUP_VMR = 0x00000009, /* setup VMR region */
GFX_CMD_ID_DESTROY_VMR = 0x0000000A, /* destroy VMR region */
GFX_CMD_ID_PROG_REG = 0x0000000B, /* program regs */
GFX_CMD_ID_GET_FW_ATTESTATION = 0x0000000F, /* Query GPUVA of the Fw Attestation DB */
/* IDs upto 0x1F are reserved for older programs (Raven, Vega 10/12/20) */
GFX_CMD_ID_LOAD_TOC = 0x00000020, /* Load TOC and obtain TMR size */
GFX_CMD_ID_AUTOLOAD_RLC = 0x00000021, /* Indicates all graphics fw loaded, start RLC autoload */
GFX_CMD_ID_BOOT_CFG = 0x00000022, /* Boot Config */
GFX_CMD_ID_SRIOV_SPATIAL_PART = 0x00000027, /* Configure spatial partitioning mode */
};
/* PSP boot config sub-commands */
enum psp_gfx_boot_config_cmd
{
BOOTCFG_CMD_SET = 1, /* Set boot configuration settings */
BOOTCFG_CMD_GET = 2, /* Get boot configuration settings */
BOOTCFG_CMD_INVALIDATE = 3 /* Reset current boot configuration settings to VBIOS defaults */
};
/* PSP boot config bitmask values */
enum psp_gfx_boot_config
{
BOOT_CONFIG_GECC = 0x1,
};
/* Command to load Trusted Application binary into PSP OS. */
struct psp_gfx_cmd_load_ta
{
unsigned int app_phy_addr_lo; /* bits [31:0] of the GPU Virtual address of the TA binary (must be 4 KB aligned) */
unsigned int app_phy_addr_hi; /* bits [63:32] of the GPU Virtual address of the TA binary */
unsigned int app_len; /* length of the TA binary in bytes */
unsigned int cmd_buf_phy_addr_lo; /* bits [31:0] of the GPU Virtual address of CMD buffer (must be 4 KB aligned) */
unsigned int cmd_buf_phy_addr_hi; /* bits [63:32] of the GPU Virtual address of CMD buffer */
unsigned int cmd_buf_len; /* length of the CMD buffer in bytes; must be multiple of 4 KB */
/* Note: CmdBufLen can be set to 0. In this case no persistent CMD buffer is provided
* for the TA. Each InvokeCommand can have dinamically mapped CMD buffer instead
* of using global persistent buffer.
*/
};
/* Command to Unload Trusted Application binary from PSP OS. */
struct psp_gfx_cmd_unload_ta
{
unsigned int session_id; /* Session ID of the loaded TA to be unloaded */
};
/* Shared buffers for InvokeCommand.
*/
struct psp_gfx_buf_desc
{
unsigned int buf_phy_addr_lo; /* bits [31:0] of GPU Virtual address of the buffer (must be 4 KB aligned) */
unsigned int buf_phy_addr_hi; /* bits [63:32] of GPU Virtual address of the buffer */
unsigned int buf_size; /* buffer size in bytes (must be multiple of 4 KB and no bigger than 64 MB) */
};
/* Max number of descriptors for one shared buffer (in how many different
* physical locations one shared buffer can be stored). If buffer is too much
* fragmented, error will be returned.
*/
#define GFX_BUF_MAX_DESC 64
struct psp_gfx_buf_list
{
unsigned int num_desc; /* number of buffer descriptors in the list */
unsigned int total_size; /* total size of all buffers in the list in bytes (must be multiple of 4 KB) */
struct psp_gfx_buf_desc buf_desc[GFX_BUF_MAX_DESC]; /* list of buffer descriptors */
/* total 776 bytes */
};
/* Command to execute InvokeCommand entry point of the TA. */
struct psp_gfx_cmd_invoke_cmd
{
unsigned int session_id; /* Session ID of the TA to be executed */
unsigned int ta_cmd_id; /* Command ID to be sent to TA */
struct psp_gfx_buf_list buf; /* one indirect buffer (scatter/gather list) */
};
/* Command to setup TMR region. */
struct psp_gfx_cmd_setup_tmr
{
unsigned int buf_phy_addr_lo; /* bits [31:0] of GPU Virtual address of TMR buffer (must be 4 KB aligned) */
unsigned int buf_phy_addr_hi; /* bits [63:32] of GPU Virtual address of TMR buffer */
unsigned int buf_size; /* buffer size in bytes (must be multiple of 4 KB) */
union {
struct {
unsigned int sriov_enabled:1; /* whether the device runs under SR-IOV*/
unsigned int virt_phy_addr:1; /* driver passes both virtual and physical address to PSP*/
unsigned int reserved:30;
} bitfield;
unsigned int tmr_flags;
};
unsigned int system_phy_addr_lo; /* bits [31:0] of system physical address of TMR buffer (must be 4 KB aligned) */
unsigned int system_phy_addr_hi; /* bits [63:32] of system physical address of TMR buffer */
};
/* FW types for GFX_CMD_ID_LOAD_IP_FW command. Limit 31. */
enum psp_gfx_fw_type {
GFX_FW_TYPE_NONE = 0, /* */
GFX_FW_TYPE_CP_ME = 1, /* CP-ME VG + RV */
GFX_FW_TYPE_CP_PFP = 2, /* CP-PFP VG + RV */
GFX_FW_TYPE_CP_CE = 3, /* CP-CE VG + RV */
GFX_FW_TYPE_CP_MEC = 4, /* CP-MEC FW VG + RV */
GFX_FW_TYPE_CP_MEC_ME1 = 5, /* CP-MEC Jump Table 1 VG + RV */
GFX_FW_TYPE_CP_MEC_ME2 = 6, /* CP-MEC Jump Table 2 VG */
GFX_FW_TYPE_RLC_V = 7, /* RLC-V VG */
GFX_FW_TYPE_RLC_G = 8, /* RLC-G VG + RV */
GFX_FW_TYPE_SDMA0 = 9, /* SDMA0 VG + RV */
GFX_FW_TYPE_SDMA1 = 10, /* SDMA1 VG */
GFX_FW_TYPE_DMCU_ERAM = 11, /* DMCU-ERAM VG + RV */
GFX_FW_TYPE_DMCU_ISR = 12, /* DMCU-ISR VG + RV */
GFX_FW_TYPE_VCN = 13, /* VCN RV */
GFX_FW_TYPE_UVD = 14, /* UVD VG */
GFX_FW_TYPE_VCE = 15, /* VCE VG */
GFX_FW_TYPE_ISP = 16, /* ISP RV */
GFX_FW_TYPE_ACP = 17, /* ACP RV */
GFX_FW_TYPE_SMU = 18, /* SMU VG */
GFX_FW_TYPE_MMSCH = 19, /* MMSCH VG */
GFX_FW_TYPE_RLC_RESTORE_LIST_GPM_MEM = 20, /* RLC GPM VG + RV */
GFX_FW_TYPE_RLC_RESTORE_LIST_SRM_MEM = 21, /* RLC SRM VG + RV */
GFX_FW_TYPE_RLC_RESTORE_LIST_SRM_CNTL = 22, /* RLC CNTL VG + RV */
GFX_FW_TYPE_UVD1 = 23, /* UVD1 VG-20 */
GFX_FW_TYPE_TOC = 24, /* TOC NV-10 */
GFX_FW_TYPE_RLC_P = 25, /* RLC P NV */
GFX_FW_TYPE_RLC_IRAM = 26, /* RLC_IRAM NV */
GFX_FW_TYPE_GLOBAL_TAP_DELAYS = 27, /* GLOBAL TAP DELAYS NV */
GFX_FW_TYPE_SE0_TAP_DELAYS = 28, /* SE0 TAP DELAYS NV */
GFX_FW_TYPE_SE1_TAP_DELAYS = 29, /* SE1 TAP DELAYS NV */
GFX_FW_TYPE_GLOBAL_SE0_SE1_SKEW_DELAYS = 30, /* GLOBAL SE0/1 SKEW DELAYS NV */
GFX_FW_TYPE_SDMA0_JT = 31, /* SDMA0 JT NV */
GFX_FW_TYPE_SDMA1_JT = 32, /* SDNA1 JT NV */
GFX_FW_TYPE_CP_MES = 33, /* CP MES NV */
GFX_FW_TYPE_MES_STACK = 34, /* MES STACK NV */
GFX_FW_TYPE_RLC_SRM_DRAM_SR = 35, /* RLC SRM DRAM NV */
GFX_FW_TYPE_RLCG_SCRATCH_SR = 36, /* RLCG SCRATCH NV */
GFX_FW_TYPE_RLCP_SCRATCH_SR = 37, /* RLCP SCRATCH NV */
GFX_FW_TYPE_RLCV_SCRATCH_SR = 38, /* RLCV SCRATCH NV */
GFX_FW_TYPE_RLX6_DRAM_SR = 39, /* RLX6 DRAM NV */
GFX_FW_TYPE_SDMA0_PG_CONTEXT = 40, /* SDMA0 PG CONTEXT NV */
GFX_FW_TYPE_SDMA1_PG_CONTEXT = 41, /* SDMA1 PG CONTEXT NV */
GFX_FW_TYPE_GLOBAL_MUX_SELECT_RAM = 42, /* GLOBAL MUX SEL RAM NV */
GFX_FW_TYPE_SE0_MUX_SELECT_RAM = 43, /* SE0 MUX SEL RAM NV */
GFX_FW_TYPE_SE1_MUX_SELECT_RAM = 44, /* SE1 MUX SEL RAM NV */
GFX_FW_TYPE_ACCUM_CTRL_RAM = 45, /* ACCUM CTRL RAM NV */
GFX_FW_TYPE_RLCP_CAM = 46, /* RLCP CAM NV */
GFX_FW_TYPE_RLC_SPP_CAM_EXT = 47, /* RLC SPP CAM EXT NV */
GFX_FW_TYPE_RLC_DRAM_BOOT = 48, /* RLC DRAM BOOT NV */
GFX_FW_TYPE_VCN0_RAM = 49, /* VCN_RAM NV + RN */
GFX_FW_TYPE_VCN1_RAM = 50, /* VCN_RAM NV + RN */
GFX_FW_TYPE_DMUB = 51, /* DMUB RN */
GFX_FW_TYPE_SDMA2 = 52, /* SDMA2 MI */
GFX_FW_TYPE_SDMA3 = 53, /* SDMA3 MI */
GFX_FW_TYPE_SDMA4 = 54, /* SDMA4 MI */
GFX_FW_TYPE_SDMA5 = 55, /* SDMA5 MI */
GFX_FW_TYPE_SDMA6 = 56, /* SDMA6 MI */
GFX_FW_TYPE_SDMA7 = 57, /* SDMA7 MI */
GFX_FW_TYPE_VCN1 = 58, /* VCN1 MI */
GFX_FW_TYPE_CAP = 62, /* CAP_FW */
GFX_FW_TYPE_SE2_TAP_DELAYS = 65, /* SE2 TAP DELAYS NV */
GFX_FW_TYPE_SE3_TAP_DELAYS = 66, /* SE3 TAP DELAYS NV */
GFX_FW_TYPE_REG_LIST = 67, /* REG_LIST MI */
GFX_FW_TYPE_IMU_I = 68, /* IMU Instruction FW SOC21 */
GFX_FW_TYPE_IMU_D = 69, /* IMU Data FW SOC21 */
GFX_FW_TYPE_LSDMA = 70, /* LSDMA FW SOC21 */
GFX_FW_TYPE_SDMA_UCODE_TH0 = 71, /* SDMA Thread 0/CTX SOC21 */
GFX_FW_TYPE_SDMA_UCODE_TH1 = 72, /* SDMA Thread 1/CTL SOC21 */
GFX_FW_TYPE_PPTABLE = 73, /* PPTABLE SOC21 */
GFX_FW_TYPE_DISCRETE_USB4 = 74, /* dUSB4 FW SOC21 */
GFX_FW_TYPE_TA = 75, /* SRIOV TA FW UUID SOC21 */
GFX_FW_TYPE_RS64_MES = 76, /* RS64 MES ucode SOC21 */
GFX_FW_TYPE_RS64_MES_STACK = 77, /* RS64 MES stack ucode SOC21 */
GFX_FW_TYPE_RS64_KIQ = 78, /* RS64 KIQ ucode SOC21 */
GFX_FW_TYPE_RS64_KIQ_STACK = 79, /* RS64 KIQ Heap stack SOC21 */
GFX_FW_TYPE_ISP_DATA = 80, /* ISP DATA SOC21 */
GFX_FW_TYPE_CP_MES_KIQ = 81, /* MES KIQ ucode SOC21 */
GFX_FW_TYPE_MES_KIQ_STACK = 82, /* MES KIQ stack SOC21 */
GFX_FW_TYPE_UMSCH_DATA = 83, /* User Mode Scheduler Data SOC21 */
GFX_FW_TYPE_UMSCH_UCODE = 84, /* User Mode Scheduler Ucode SOC21 */
GFX_FW_TYPE_UMSCH_CMD_BUFFER = 85, /* User Mode Scheduler Command Buffer SOC21 */
GFX_FW_TYPE_USB_DP_COMBO_PHY = 86, /* USB-Display port Combo SOC21 */
GFX_FW_TYPE_RS64_PFP = 87, /* RS64 PFP SOC21 */
GFX_FW_TYPE_RS64_ME = 88, /* RS64 ME SOC21 */
GFX_FW_TYPE_RS64_MEC = 89, /* RS64 MEC SOC21 */
GFX_FW_TYPE_RS64_PFP_P0_STACK = 90, /* RS64 PFP stack P0 SOC21 */
GFX_FW_TYPE_RS64_PFP_P1_STACK = 91, /* RS64 PFP stack P1 SOC21 */
GFX_FW_TYPE_RS64_ME_P0_STACK = 92, /* RS64 ME stack P0 SOC21 */
GFX_FW_TYPE_RS64_ME_P1_STACK = 93, /* RS64 ME stack P1 SOC21 */
GFX_FW_TYPE_RS64_MEC_P0_STACK = 94, /* RS64 MEC stack P0 SOC21 */
GFX_FW_TYPE_RS64_MEC_P1_STACK = 95, /* RS64 MEC stack P1 SOC21 */
GFX_FW_TYPE_RS64_MEC_P2_STACK = 96, /* RS64 MEC stack P2 SOC21 */
GFX_FW_TYPE_RS64_MEC_P3_STACK = 97, /* RS64 MEC stack P3 SOC21 */
GFX_FW_TYPE_VPEC_FW1 = 100, /* VPEC FW1 To Save VPE */
GFX_FW_TYPE_VPEC_FW2 = 101, /* VPEC FW2 To Save VPE */
GFX_FW_TYPE_VPE = 102,
GFX_FW_TYPE_JPEG_RAM = 128, /**< JPEG Command buffer */
GFX_FW_TYPE_P2S_TABLE = 129,
GFX_FW_TYPE_MAX
};
/* Command to load HW IP FW. */
struct psp_gfx_cmd_load_ip_fw
{
unsigned int fw_phy_addr_lo; /* bits [31:0] of GPU Virtual address of FW location (must be 4 KB aligned) */
unsigned int fw_phy_addr_hi; /* bits [63:32] of GPU Virtual address of FW location */
unsigned int fw_size; /* FW buffer size in bytes */
enum psp_gfx_fw_type fw_type; /* FW type */
};
/* Command to save/restore HW IP FW. */
struct psp_gfx_cmd_save_restore_ip_fw
{
unsigned int save_fw; /* if set, command is used for saving fw otherwise for resetoring*/
unsigned int save_restore_addr_lo; /* bits [31:0] of FB address of GART memory used as save/restore buffer (must be 4 KB aligned) */
unsigned int save_restore_addr_hi; /* bits [63:32] of FB address of GART memory used as save/restore buffer */
unsigned int buf_size; /* Size of the save/restore buffer in bytes */
enum psp_gfx_fw_type fw_type; /* FW type */
};
/* Command to setup register program */
struct psp_gfx_cmd_reg_prog {
unsigned int reg_value;
unsigned int reg_id;
};
/* Command to load TOC */
struct psp_gfx_cmd_load_toc
{
unsigned int toc_phy_addr_lo; /* bits [31:0] of GPU Virtual address of FW location (must be 4 KB aligned) */
unsigned int toc_phy_addr_hi; /* bits [63:32] of GPU Virtual address of FW location */
unsigned int toc_size; /* FW buffer size in bytes */
};
/* Dynamic boot configuration */
struct psp_gfx_cmd_boot_cfg
{
unsigned int timestamp; /* calendar time as number of seconds */
enum psp_gfx_boot_config_cmd sub_cmd; /* sub-command indicating how to process command data */
unsigned int boot_config; /* dynamic boot configuration bitmask */
unsigned int boot_config_valid; /* dynamic boot configuration valid bits bitmask */
};
struct psp_gfx_cmd_sriov_spatial_part {
unsigned int mode;
unsigned int override_ips;
unsigned int override_xcds_avail;
unsigned int override_this_aid;
};
/* All GFX ring buffer commands. */
union psp_gfx_commands
{
struct psp_gfx_cmd_load_ta cmd_load_ta;
struct psp_gfx_cmd_unload_ta cmd_unload_ta;
struct psp_gfx_cmd_invoke_cmd cmd_invoke_cmd;
struct psp_gfx_cmd_setup_tmr cmd_setup_tmr;
struct psp_gfx_cmd_load_ip_fw cmd_load_ip_fw;
struct psp_gfx_cmd_save_restore_ip_fw cmd_save_restore_ip_fw;
struct psp_gfx_cmd_reg_prog cmd_setup_reg_prog;
struct psp_gfx_cmd_setup_tmr cmd_setup_vmr;
struct psp_gfx_cmd_load_toc cmd_load_toc;
struct psp_gfx_cmd_boot_cfg boot_cfg;
struct psp_gfx_cmd_sriov_spatial_part cmd_spatial_part;
};
struct psp_gfx_uresp_reserved
{
unsigned int reserved[8];
};
/* Command-specific response for Fw Attestation Db */
struct psp_gfx_uresp_fwar_db_info
{
unsigned int fwar_db_addr_lo;
unsigned int fwar_db_addr_hi;
};
/* Command-specific response for boot config. */
struct psp_gfx_uresp_bootcfg {
unsigned int boot_cfg; /* boot config data */
};
/* Union of command-specific responses for GPCOM ring. */
union psp_gfx_uresp {
struct psp_gfx_uresp_reserved reserved;
struct psp_gfx_uresp_bootcfg boot_cfg;
struct psp_gfx_uresp_fwar_db_info fwar_db_info;
};
/* Structure of GFX Response buffer.
* For GPCOM I/F it is part of GFX_CMD_RESP buffer, for RBI
* it is separate buffer.
*/
struct psp_gfx_resp
{
unsigned int status; /* +0 status of command execution */
unsigned int session_id; /* +4 session ID in response to LoadTa command */
unsigned int fw_addr_lo; /* +8 bits [31:0] of FW address within TMR (in response to cmd_load_ip_fw command) */
unsigned int fw_addr_hi; /* +12 bits [63:32] of FW address within TMR (in response to cmd_load_ip_fw command) */
unsigned int tmr_size; /* +16 size of the TMR to be reserved including MM fw and Gfx fw in response to cmd_load_toc command */
unsigned int reserved[11];
union psp_gfx_uresp uresp; /* +64 response union containing command-specific responses */
/* total 96 bytes */
};
/* Structure of Command buffer pointed by psp_gfx_rb_frame.cmd_buf_addr_hi
* and psp_gfx_rb_frame.cmd_buf_addr_lo.
*/
struct psp_gfx_cmd_resp
{
unsigned int buf_size; /* +0 total size of the buffer in bytes */
unsigned int buf_version; /* +4 version of the buffer strusture; must be PSP_GFX_CMD_BUF_VERSION */
unsigned int cmd_id; /* +8 command ID */
/* These fields are used for RBI only. They are all 0 in GPCOM commands
*/
unsigned int resp_buf_addr_lo; /* +12 bits [31:0] of GPU Virtual address of response buffer (must be 4 KB aligned) */
unsigned int resp_buf_addr_hi; /* +16 bits [63:32] of GPU Virtual address of response buffer */
unsigned int resp_offset; /* +20 offset within response buffer */
unsigned int resp_buf_size; /* +24 total size of the response buffer in bytes */
union psp_gfx_commands cmd; /* +28 command specific structures */
unsigned char reserved_1[864 - sizeof(union psp_gfx_commands) - 28];
/* Note: Resp is part of this buffer for GPCOM ring. For RBI ring the response
* is separate buffer pointed by resp_buf_addr_hi and resp_buf_addr_lo.
*/
struct psp_gfx_resp resp; /* +864 response */
unsigned char reserved_2[1024 - 864 - sizeof(struct psp_gfx_resp)];
/* total size 1024 bytes */
};
#define FRAME_TYPE_DESTROY 1 /* frame sent by KMD driver when UMD Scheduler context is destroyed*/
/* Structure of the Ring Buffer Frame */
struct psp_gfx_rb_frame
{
unsigned int cmd_buf_addr_lo; /* +0 bits [31:0] of GPU Virtual address of command buffer (must be 4 KB aligned) */
unsigned int cmd_buf_addr_hi; /* +4 bits [63:32] of GPU Virtual address of command buffer */
unsigned int cmd_buf_size; /* +8 command buffer size in bytes */
unsigned int fence_addr_lo; /* +12 bits [31:0] of GPU Virtual address of Fence for this frame */
unsigned int fence_addr_hi; /* +16 bits [63:32] of GPU Virtual address of Fence for this frame */
unsigned int fence_value; /* +20 Fence value */
unsigned int sid_lo; /* +24 bits [31:0] of SID value (used only for RBI frames) */
unsigned int sid_hi; /* +28 bits [63:32] of SID value (used only for RBI frames) */
unsigned char vmid; /* +32 VMID value used for mapping of all addresses for this frame */
unsigned char frame_type; /* +33 1: destory context frame, 0: all other frames; used only for RBI frames */
unsigned char reserved1[2]; /* +34 reserved, must be 0 */
unsigned int reserved2[7]; /* +36 reserved, must be 0 */
/* total 64 bytes */
};
#define PSP_ERR_UNKNOWN_COMMAND 0x00000100
enum tee_error_code {
TEE_SUCCESS = 0x00000000,
TEE_ERROR_NOT_SUPPORTED = 0xFFFF000A,
};
#endif /* _PSP_TEE_GFX_IF_H_ */

1194
artifacts/package_sources/tinygrad/extra/amdpci/headers/v11_structs.h Normal file
View File

File diff suppressed because it is too large Load Diff

1191
artifacts/package_sources/tinygrad/extra/amdpci/headers/v12_structs.h Normal file
View File

File diff suppressed because it is too large Load Diff

20
artifacts/package_sources/tinygrad/extra/amdpci/hive_reset.py Executable file
View File

@@ -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()

7
artifacts/package_sources/tinygrad/extra/amdpci/overlay/gc_9_4_3.h Normal file
View File

@@ -0,0 +1,7 @@
// From MQD struct
#define regCOMPUTE_CURRENT_LOGIC_XCC_ID 0x0e25
#define regCOMPUTE_CURRENT_LOGIC_XCC_ID_BASE_IDX 0
// Mask is probably not full register, doesn't matter though
#define COMPUTE_CURRENT_LOGIC_XCC_ID__CURRENT_LOGIC_XCC_ID__SHIFT 0x0
#define COMPUTE_CURRENT_LOGIC_XCC_ID__CURRENT_LOGIC_XCC_ID_MASK 0xFFFFFFFFL

96
artifacts/package_sources/tinygrad/extra/amdpci/proclogs.py Normal file
View File

@@ -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()

3
artifacts/package_sources/tinygrad/extra/amdpci/setup_python_cap.sh Executable file
View File

@@ -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

2
artifacts/package_sources/tinygrad/extra/amdpci/setup_vfio.sh Executable file
View File

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

154
artifacts/package_sources/tinygrad/extra/archprobe.py Normal file
View File

@@ -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()

111
artifacts/package_sources/tinygrad/extra/bench_log.py Normal file
View File

@@ -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)

31
artifacts/package_sources/tinygrad/extra/benchmark_llm.py Normal file
View File

@@ -0,0 +1,31 @@
import argparse, time
from tinygrad.llm.model import Transformer
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("--model", required=True, help="path to gguf model")
parser.add_argument("--max-context", type=int, default=8192, help="max context length (default: %(default)s)")
parser.add_argument("--prompt-tokens", type=int, default=1024, help="number of prompt tokens (default: %(default)s)")
parser.add_argument("--decode-tokens", type=int, default=16, help="number of tokens to decode (default: %(default)s)")
parser.add_argument("--chunk-size", type=int, default=32, help="chunk size for prefill (default: %(default)s)")
args = parser.parse_args()
st = time.perf_counter()
model, _ = Transformer.from_gguf(args.model, args.max_context)
print(f"load {time.perf_counter()-st:.3f}s", flush=True)
st = time.perf_counter()
model.warmup()
print(f"warm {time.perf_counter()-st:.3f}s", flush=True)
prompt = [257] + [1000+i%1000 for i in range(args.prompt_tokens-1)]
gen = model.generate(prompt, chunk_size=args.chunk_size)
st = time.perf_counter()
# first token is time-to-first-token; counted as part of prefill
output = [next(gen)]
pt = time.perf_counter()
print(f"prefill {args.prompt_tokens/(pt-st):.3f} tok/s", flush=True)
for _ in range(args.decode_tokens): output.append(next(gen))
et = time.perf_counter()
print(f"decode {args.decode_tokens/(et-pt):.3f} tok/s output {output}", flush=True)

4
artifacts/package_sources/tinygrad/extra/cl_android.sh Normal file
View File

@@ -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

4
artifacts/package_sources/tinygrad/extra/datasets/.gitignore vendored Normal file
View File

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

43
artifacts/package_sources/tinygrad/extra/datasets/__init__.py Normal file
View File

@@ -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

41
artifacts/package_sources/tinygrad/extra/datasets/fake_imagenet_from_mnist.py Executable file
View File

@@ -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"))

91
artifacts/package_sources/tinygrad/extra/datasets/imagenet.py Normal file
View File

@@ -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

51
artifacts/package_sources/tinygrad/extra/datasets/imagenet_download.py Normal file
View File

@@ -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()

219
artifacts/package_sources/tinygrad/extra/datasets/kits19.py Normal file
View File

@@ -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)

82
artifacts/package_sources/tinygrad/extra/datasets/librispeech.py Normal file
View File

@@ -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)

209
artifacts/package_sources/tinygrad/extra/datasets/openimages.py Normal file
View File

@@ -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")

21
artifacts/package_sources/tinygrad/extra/datasets/preprocess_imagenet.py Normal file
View File

@@ -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

BIN
artifacts/package_sources/tinygrad/extra/datasets/sops.gz Normal file
View File

Binary file not shown.

148
artifacts/package_sources/tinygrad/extra/datasets/squad.py Normal file
View File

@@ -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)

398
artifacts/package_sources/tinygrad/extra/datasets/wikipedia.py Normal file
View File

@@ -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))

54
artifacts/package_sources/tinygrad/extra/datasets/wikipedia_download.py Normal file
View File

@@ -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")))

38
artifacts/package_sources/tinygrad/extra/dsp/Dockerfile Normal file
View File

@@ -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"]

125
artifacts/package_sources/tinygrad/extra/dsp/compile.py Executable file
View File

@@ -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)

3
artifacts/package_sources/tinygrad/extra/dsp/gen.sh Executable file
View File

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

101
artifacts/package_sources/tinygrad/extra/dsp/hook.py Normal file
View File

@@ -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)

1146
artifacts/package_sources/tinygrad/extra/dsp/include/AEEStdDef.h Normal file
View File

File diff suppressed because it is too large Load Diff

328
artifacts/package_sources/tinygrad/extra/dsp/include/AEEStdErr.h Executable file
View File

@@ -0,0 +1,328 @@
/*
* Copyright (c) 2005-2007, 2012-2013, 2019-2020 Qualcomm Technologies, Inc.
* All Rights Reserved.
* Redistribution and use in source and binary forms, with or without
* modification, are permitted provided that the following conditions are met:
*
* 1. Redistributions of source code must retain the above copyright notice,
* this list of conditions and the following disclaimer.
*
* 2. Redistributions in binary form must reproduce the above copyright notice,
* this list of conditions and the following disclaimer in the documentation
* and/or other materials provided with the distribution.
*
* 3. Neither the name of the copyright holder nor the names of its contributors
* may be used to endorse or promote products derived from this software without
* specific prior written permission.
*
* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
* AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
* IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
* ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
* LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
* CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
* SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
* INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
* CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
* ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
* POSSIBILITY OF SUCH DAMAGE.
*/
#ifndef AEESTDERR_H
#define AEESTDERR_H
//
// Basic Error Codes
//
//
#if defined(__hexagon__)
#define AEE_EOFFSET 0x80000400
#else
#define AEE_EOFFSET 0x00000000
#endif
/** @defgroup stdbasicerror Basic error codes
* @{
*/
#define AEE_SUCCESS 0 ///< No error
#define AEE_EUNKNOWN -1 ///< Unknown error (should not use this)
#define AEE_EFAILED (AEE_EOFFSET + 0x001) ///< General failure
#define AEE_ENOMEMORY (AEE_EOFFSET + 0x002) ///< Memory allocation failed because of insufficient RAM
#define AEE_ECLASSNOTSUPPORT (AEE_EOFFSET + 0x003) ///< Specified class unsupported
#define AEE_EVERSIONNOTSUPPORT (AEE_EOFFSET + 0x004) ///< Version not supported
#define AEE_EALREADYLOADED (AEE_EOFFSET + 0x005) ///< Object already loaded
#define AEE_EUNABLETOLOAD (AEE_EOFFSET + 0x006) ///< Unable to load object/applet
#define AEE_EUNABLETOUNLOAD (AEE_EOFFSET + 0x007) ///< Unable to unload
///< object/applet
#define AEE_EALARMPENDING (AEE_EOFFSET + 0x008) ///< Alarm is pending
#define AEE_EINVALIDTIME (AEE_EOFFSET + 0x009) ///< Invalid time
#define AEE_EBADCLASS (AEE_EOFFSET + 0x00A) ///< NULL class object
#define AEE_EBADMETRIC (AEE_EOFFSET + 0x00B) ///< Invalid metric specified
#define AEE_EEXPIRED (AEE_EOFFSET + 0x00C) ///< App/Component Expired
#define AEE_EBADSTATE (AEE_EOFFSET + 0x00D) ///< Process or thread is not in expected state
#define AEE_EBADPARM (AEE_EOFFSET + 0x00E) ///< Invalid parameter
#define AEE_ESCHEMENOTSUPPORTED (AEE_EOFFSET + 0x00F) ///< Invalid URL scheme
#define AEE_EBADITEM (AEE_EOFFSET + 0x010) ///< Value out of range
#define AEE_EINVALIDFORMAT (AEE_EOFFSET + 0x011) ///< Invalid format
#define AEE_EINCOMPLETEITEM (AEE_EOFFSET + 0x012) ///< Incomplete item, like length of a string is less that expected
#define AEE_ENOPERSISTMEMORY (AEE_EOFFSET + 0x013) ///< Insufficient flash
#define AEE_EUNSUPPORTED (AEE_EOFFSET + 0x014) ///< API not implemented
#define AEE_EPRIVLEVEL (AEE_EOFFSET + 0x015) ///< Privileges are insufficient
///< for this operation
#define AEE_ERESOURCENOTFOUND (AEE_EOFFSET + 0x016) ///< Unable to find specified
///< resource
#define AEE_EREENTERED (AEE_EOFFSET + 0x017) ///< Non re-entrant API
///< re-entered
#define AEE_EBADTASK (AEE_EOFFSET + 0x018) ///< API called in wrong task
///< context
#define AEE_EALLOCATED (AEE_EOFFSET + 0x019) ///< App/Module left memory
///< allocated when released.
#define AEE_EALREADY (AEE_EOFFSET + 0x01A) ///< Operation is already in
///< progress
#define AEE_EADSAUTHBAD (AEE_EOFFSET + 0x01B) ///< ADS mutual authorization
///< failed
#define AEE_ENEEDSERVICEPROG (AEE_EOFFSET + 0x01C) ///< Need service programming
#define AEE_EMEMPTR (AEE_EOFFSET + 0x01D) ///< bad memory pointer, expected to be NULL
#define AEE_EHEAP (AEE_EOFFSET + 0x01E) ///< An internal heap error was detected
#define AEE_EIDLE (AEE_EOFFSET + 0x01F) ///< Context (system, interface,
///< etc.) is idle
#define AEE_EITEMBUSY (AEE_EOFFSET + 0x020) ///< Context (system, interface,
///< etc.) is busy
#define AEE_EBADSID (AEE_EOFFSET + 0x021) ///< Invalid subscriber ID
#define AEE_ENOTYPE (AEE_EOFFSET + 0x022) ///< No type detected/found
#define AEE_ENEEDMORE (AEE_EOFFSET + 0x023) ///< Need more data/info
#define AEE_EADSCAPS (AEE_EOFFSET + 0x024) ///< ADS Capabilities do not
///< match those required for phone
#define AEE_EBADSHUTDOWN (AEE_EOFFSET + 0x025) ///< App failed to close properly
#define AEE_EBUFFERTOOSMALL (AEE_EOFFSET + 0x026) ///< Destination buffer given is
///< too small
///< or service exists or is
///< valid
#define AEE_EACKPENDING (AEE_EOFFSET + 0x028) ///< ACK pending on application
#define AEE_ENOTOWNER (AEE_EOFFSET + 0x029) ///< Not an owner authorized to
///< perform the operation
#define AEE_EINVALIDITEM (AEE_EOFFSET + 0x02A) ///< Current item is invalid, it can be a switch case or a pointer to memory
#define AEE_ENOTALLOWED (AEE_EOFFSET + 0x02B) ///< Not allowed to perform the
///< operation
#define AEE_EBADHANDLE (AEE_EOFFSET + 0x02C) ///< Invalid/Wrong handle
#define AEE_EINVHANDLE (AEE_EOFFSET + 0x02C) ///< Invalid handle - adding here as its defined in vendor AEEStdErr.h - needed to check valid handle in stub.c
#define AEE_EOUTOFHANDLES (AEE_EOFFSET + 0x02D) ///< Out of handles (Handle list is already full)
//Hole here
#define AEE_ENOMORE (AEE_EOFFSET + 0x02F) ///< No more items available --
///< reached end
#define AEE_ECPUEXCEPTION (AEE_EOFFSET + 0x030) ///< A CPU exception occurred
#define AEE_EREADONLY (AEE_EOFFSET + 0x031) ///< Cannot change read-only
///< object or parameter ( Parameter is in protected mode)
#define AEE_ERPC (AEE_EOFFSET + 0x200) ///< Error due to fastrpc implementation
#define AEE_EFILE (AEE_EOFFSET + 0x201) ///<File handling related error
//NOTE: Used in both HLOS and DSP.
#define AEE_ENOSUCH (39) ///< No such name, port, socket
#define AEE_EINTERRUPTED (46) ///< Waitable call is interrupted,
///< the user should return to the HLOS and retry the call
#define AEE_ECONNRESET (104) ///< Connection reset by peer
#define AEE_EWOULDBLOCK (516) ///< Operation would block if not
///< non-blocking; wait and try
///< again
/**
* @}
*/
/** @defgroup sigverifyerror Sigverify error codes
* @{
*/
#define AEE_EINVALIDMSG (AEE_EOFFSET + 0x032) ///< Invalid SMD message from APPS
#define AEE_EINVALIDTHREAD (AEE_EOFFSET + 0x033) ///< Invalid thread
#define AEE_EINVALIDPROCESS (AEE_EOFFSET + 0x034) ///< Invalid Process
#define AEE_EINVALIDFILENAME (AEE_EOFFSET + 0x035) ///< Invalid filename
#define AEE_EINVALIDDIGESTSIZE (AEE_EOFFSET + 0x036) ///< Invalid digest size
#define AEE_EINVALIDSEGS (AEE_EOFFSET + 0x037) ///< Invalid segments
#define AEE_EINVALIDSIGNATURE (AEE_EOFFSET + 0x038) ///< Invalid signature
#define AEE_EINVALIDDOMAIN (AEE_EOFFSET + 0x039) ///< Invalid DSP domain
#define AEE_EINVALIDFD (AEE_EOFFSET + 0x03A) ///< Invalid file descriptor
#define AEE_EINVALIDDEVICE (AEE_EOFFSET + 0x03B) ///< Invalid Device or Device node open failed for the domain
#define AEE_EINVALIDMODE (AEE_EOFFSET + 0x03C) ///< Invalid Mode
#define AEE_EINVALIDPROCNAME (AEE_EOFFSET + 0x03D) ///< Invalid Process name
#define AEE_ENOSUCHMOD (AEE_EOFFSET + 0x03E) ///< No such module
#define AEE_ENOSUCHINSTANCE (AEE_EOFFSET + 0x03F) ///< No instance in the list lookup
#define AEE_ENOSUCHTHREAD (AEE_EOFFSET + 0x040) ///< No such thread
#define AEE_ENOSUCHPROCESS (AEE_EOFFSET + 0x041) ///< No such process
#define AEE_ENOSUCHSYMBOL (AEE_EOFFSET + 0x042) ///< No such symbol( dlsym for the symbol failed)
#define AEE_ENOSUCHDEVICE (AEE_EOFFSET + 0x043) ///< No such device
#define AEE_ENOSUCHPROP (AEE_EOFFSET + 0x044) ///< No such dal property
#define AEE_ENOSUCHFILE (AEE_EOFFSET + 0x045) ///< No such file found
#define AEE_ENOSUCHHANDLE (AEE_EOFFSET + 0x046) ///< No such handle
#define AEE_ENOSUCHSTREAM (AEE_EOFFSET + 0x047) ///< No such stream
#define AEE_ENOSUCHMAP (AEE_EOFFSET + 0x048) ///< No mapping exists for this address on DSP
#define AEE_ENOSUCHREGISTER (AEE_EOFFSET + 0x049) ///< No such register
#define AEE_ENOSUCHCLIENT (AEE_EOFFSET + 0x04A) ///< No such QDI client
#define AEE_EBADDOMAIN (AEE_EOFFSET + 0x04B) ///< Bad domain (not initialized)
#define AEE_EBADOFFSET (AEE_EOFFSET + 0x04C) ///< Bad buffer/page/heap offset
#define AEE_EBADSIZE (AEE_EOFFSET + 0x04D) ///< Bad buffer/page/heap size
#define AEE_EBADPERMS (AEE_EOFFSET + 0x04E) ///< Bad FILE/MAP/MEM permissions
#define AEE_EBADFD (AEE_EOFFSET + 0x04F) ///< Bad file descriptor
#define AEE_EBADPID (AEE_EOFFSET + 0x050) ///< Bad PID from HLOS
#define AEE_EBADTID (AEE_EOFFSET + 0x051) ///< Bad TID
#define AEE_EBADELF (AEE_EOFFSET + 0x052) ///< Bad elf file
#define AEE_EBADASID (AEE_EOFFSET + 0x053) ///< Bad asid
#define AEE_EBADCONTEXT (AEE_EOFFSET + 0x054) ///< Bad context
#define AEE_EBADMEMALIGN (AEE_EOFFSET + 0x055) ///< Bad memory alignment
#define AEE_EIOCTL (AEE_EOFFSET + 0x056) ///< ioctl call failed
#define AEE_EFOPEN (AEE_EOFFSET + 0x057) ///< file open error or device node open failed for DSP domain
#define AEE_EFGETS (AEE_EOFFSET + 0x058) ///< file get string error
#define AEE_EFFLUSH (AEE_EOFFSET + 0x059) ///< file flush error
#define AEE_EFCLOSE (AEE_EOFFSET + 0x05A) ///< file close error
#define AEE_EEOF (AEE_EOFFSET + 0x05B) ///< File EOF reached
#define AEE_EFREAD (AEE_EOFFSET + 0x05C) ///< file read failed
#define AEE_EFWRITE (AEE_EOFFSET + 0x05D) ///< file write failed
#define AEE_EFGETPOS (AEE_EOFFSET + 0x05E) ///< file get position failed
#define AEE_EFSETPOS (AEE_EOFFSET + 0x05F) ///< file set position failed
#define AEE_EFTELL (AEE_EOFFSET + 0x060) ///< file tell position failed
#define AEE_EFSEEK (AEE_EOFFSET + 0x061) ///< file seek failed
#define AEE_EFLEN (AEE_EOFFSET + 0x062) ///< file len greater than expected
#define AEE_EGETENV (AEE_EOFFSET + 0x063) ///< apps_std get enviroment failed
#define AEE_ESETENV (AEE_EOFFSET + 0x064) ///< apps_std set enviroment failed
#define AEE_EMMAP (AEE_EOFFSET + 0x065) ///< mmap failed
#define AEE_EIONMAP (AEE_EOFFSET + 0x066) ///< ion map failed
#define AEE_EIONALLOC (AEE_EOFFSET + 0x067) ///< ion alloc failed
#define AEE_ENORPCMEMORY (AEE_EOFFSET + 0x068) ///< ION memory allocation failed
#define AEE_ENOROOTOFTRUST (AEE_EOFFSET + 0x069) ///< No root of trust for sigverify
#define AEE_ENOTLOCKED (AEE_EOFFSET + 0x06A) ///< Unlock failed, not locked before
#define AEE_ENOTINITIALIZED (AEE_EOFFSET + 0x06B) ///< Not initialized
#define AEE_EUNSUPPORTEDAPI (AEE_EOFFSET + 0x06C) ///< unsupported API/request ID
#define AEE_EUNPACK (AEE_EOFFSET + 0x06D) ///< unpacking command failed
#define AEE_EPOLL (AEE_EOFFSET + 0x06E) ///< error while polling for event
#define AEE_EEVENTREAD (AEE_EOFFSET + 0x06F) ///< event read failed
#define AEE_EMAXBUFS (AEE_EOFFSET + 0x070) ///< Maximum buffers
#define AEE_EINVARGS (AEE_EOFFSET + 0x071) ///< Invalid Arguments
#define AEE_ECONNREFUSED (AEE_EOFFSET + 0x072) ///< Connection refused to DSP
#define AEE_EUNSIGNEDMOD (AEE_EOFFSET + 0x081) ///< test-sig not found, Unsigned shared object
#define AEE_EINVALIDHASH (AEE_EOFFSET + 0x082) ///< test-sig not found, Invalid hash object
#define AEE_EBADVA (AEE_EOFFSET + 0x083) ///< Bad VA address
#define AEE_ENOSUCHJOB (AEE_EOFFSET + 0x084) ///< No such job
#define AEE_ENOSUCHGROUP (AEE_EOFFSET + 0x084) ///< No such static pd group
#define AEE_EBADMAPREFCNT (AEE_EOFFSET + 0x085) ///< Bad map reference count
#define AEE_EBADPAGECNT (AEE_EOFFSET + 0x086) ///< Bad page count
#define AEE_EMAPALREADYPRESENT (AEE_EOFFSET + 0x087) ///< Map already present
#define AEE_ENOFREESECTION (AEE_EOFFSET + 0x088) ///< No more free sections available
#define AEE_U2GCLIENT_OPEN (AEE_EOFFSET + 0x089) ///< u2g client open failed
/**
* @}
*/
/** @defgroup smderror SMD error codes
* @{
*/
#if defined(__hexagon__)
#define AEE_EGLINK_OFFSET (AEE_EOFFSET + 0x100) ///< SMD errors offset
#define AEE_EGLINKBADPACKET (AEE_EOFFSET + 0x101) ///< SMD invalid packet size
#define AEE_EGLINKALREADYOPEN (AEE_EOFFSET + 0x102) ///< SMD port is already open
#define AEE_EGLINKOPENFAILED (AEE_EOFFSET + 0x103) ///< SMD port open failed
#define AEE_EGLINKWRITE (AEE_EOFFSET + 0x104) ///< SMD port write failed
#define AEE_EGLINKREGISTER (AEE_EOFFSET + 0x105) ///< SMD port register callback failed
#else
#define AEE_ESMD_OFFSET (AEE_EOFFSET + 0x100) ///< SMD errors offset
#define AEE_ESMDBADPACKET (AEE_EOFFSET + 0x101) ///< SMD invalid packet size
#define AEE_ESMDALREADYOPEN (AEE_EOFFSET + 0x102) ///< SMD port is already open
#define AEE_ESMDOPENFAILED (AEE_EOFFSET + 0x103) ///< SMD port open failed
#endif
/**
* @}
*/
/** @defgroup dalerror DAL error codes
* @{
*/
#define AEE_EDAL_OFFSET (AEE_EOFFSET + 0x120) ///< Dal error offset
#define AEE_EDALDEVATTACH (AEE_EOFFSET + 0x121) ///< DAL attach error
#define AEE_EDALINTREGISTER (AEE_EOFFSET + 0x122) ///< DAL interrupt register error
#define AEE_EDALINTUNREGISTER (AEE_EOFFSET + 0x123) ///< Dal interrupt unregister error
#define AEE_EDALGETPROP (AEE_EOFFSET + 0x124) ///< Dal get property
#define AEE_EDALGETVAL (AEE_EOFFSET + 0x125) ///< Dal get property value
#define AEE_EDCVSREQUEST (AEE_EOFFSET + 0x126) ///< Dal get property value
/**
* @}
*/
/** @defgroup qurterror QURT error codes
* @{
*/
#define AEE_EQURT_OFFSET (AEE_EOFFSET + 0x140) ///< QURT error offset
#define AEE_EQURTREGIONCREATE (AEE_EOFFSET + 0x141) ///< QURT region create failed
#define AEE_EQURTCACHECLEAN (AEE_EOFFSET + 0x142) ///< QURT cache clean failed
#define AEE_EQURTREGIONGETATTR (AEE_EOFFSET + 0x143) ///< QURT region get attribute failed
#define AEE_EQURTBADREGIONPERMS (AEE_EOFFSET + 0x144) ///< QURT bad permissions for region
#define AEE_EQURTMEMPOOLADD (AEE_EOFFSET + 0x145) ///< QURT Add to memory pool failed
#define AEE_EQURTREGISTERDEV (AEE_EOFFSET + 0x146) ///< QURT register device failed
#define AEE_EQURTMEMPOOLCREATE (AEE_EOFFSET + 0x147) ///< QURT create memory pool failed
#define AEE_EQURTGETVA (AEE_EOFFSET + 0x148) ///< QURT get VA failed
#define AEE_EQURTREGIONDELETE (AEE_EOFFSET + 0x149) ///< QURT region delete failed
#define AEE_EQURTMEMPOOLATTACH (AEE_EOFFSET + 0x14A) ///< QURT memory pool attach failed
#define AEE_EQURTTHREADCREATE (AEE_EOFFSET + 0x14B) ///< QURT thread create failed
#define AEE_EQURTCOPYTOUSER (AEE_EOFFSET + 0x14C) ///< QURT copy to user memory failed
#define AEE_EQURTMEMMAPCREATE (AEE_EOFFSET + 0x14D) ///< QURT map create failed
#define AEE_EQURTINVHANDLE (AEE_EOFFSET + 0x14E) ///< QURT Invalid client handle
#define AEE_EQURTBADASID (AEE_EOFFSET + 0x14F) ///< QURT Bad ASIC from QURT
#define AEE_EQURTOPENFAILED (AEE_EOFFSET + 0x150) ///< QURT QDI open failed
#define AEE_EQURTCOPYFROMUSER (AEE_EOFFSET + 0x151) ///< QURT Copy from user failed
#define AEE_EQURTLINELOCK (AEE_EOFFSET + 0x152) ///< QURT Line lock failed
#define AEE_EQURTQDIDEFMETHOD (AEE_EOFFSET + 0x153) ///< QURT QDI default method failed
#define AEE_EQURTCREATEHANDLE (AEE_EOFFSET + 0x154) ///< QURT create handle from obj failed
#define AEE_EQURTWRITABLEMEM (AEE_EOFFSET + 0x155) ///< QURT CPZ migration writable mem
#define AEE_EQURTTHREADCREATEDEF (AEE_EOFFSET + 0x156) ///< QURT thread create def
#define AEE_EQURTLOOKUPVA (AEE_EOFFSET + 0x157) ///< QURT lookup VA
#define AEE_EQURTLOOKUPPA (AEE_EOFFSET + 0x158) ///< QURT lookup PA
#define AEE_EQURTMIGRATESECURE (AEE_EOFFSET + 0x159) ///< QURT CPZ migration failure
#define AEE_EQURTQDIOPEN (AEE_EOFFSET + 0X160) ///< QURT QDI open failure
#define AEE_EQURTMAPREMOVE (AEE_EOFFSET + 0X161) ///< QURT map remove failure
#define AEE_EQURTQDICLOSE (AEE_EOFFSET + 0X162) ///< QURT QDI close failed
#define AEE_EQURTWAIT (AEE_EOFFSET + 0X163) ///< QURT Futex wait failed
/**
* @}
*/
/** @defgroup mmpmerr MMPM error codes
* @{
*/
#define AEE_EMMPM_OFFSET (AEE_EOFFSET + 0x170) ///< MMPM errors offset
#define AEE_EMMPMREQUEST (AEE_EOFFSET + 0x171) ///< MMPM Power request to failed
#define AEE_EMMPMRELEASE (AEE_EOFFSET + 0x172) ///< MMPM Release request failed
#define AEE_EMMPMSETPARAM (AEE_EOFFSET + 0x173) ///< MMPM set param request failed
#define AEE_EMMPMREGISTER (AEE_EOFFSET + 0x174) ///< MMPM Register request failed
#define AEE_EMMPMGETINFO (AEE_EOFFSET + 0x175) ///< MMPM Get info failed
#define AEE_EMAX_MMPM_CLIENTS (AEE_EOFFSET + 0x176) ///< MMPM Reached maximum clients per PD(HAP_MAX_CLIENTS)
#define AEE_EDCVSREGISTER (AEE_EOFFSET + 0x177) ///< ADSP DCVS client registration failed
#define AEE_PDRREGFAIL (AEE_EOFFSET + 0x178) ///< Error Callback Services Registration failed for PD
/**
* @}
*/
#define AEE_DEFAULT_PROCESS (AEE_EOFFSET + 0x180) ///< Default process in Guest OS is not present
#define AEE_ENULLCONTEXT (AEE_EOFFSET + 0x181) ///< User NULL context vote
#define AEE_EINVALIDJOB (AEE_EOFFSET + 0x182) ///< AsyncRPC Invalid job
#define AEE_EBUSY (AEE_EOFFSET + 0x183) ///< AsyncRPC Pending job
/** @defgroup heaperror Heap error codes
* @{
*/
#define E_APPS_BUSY_RETRY_LATER (AEE_EOFFSET + 0x190) ///< Retry because the apps is busy
#define E_HLOS_CAP_REACHED (AEE_EOFFSET + 0x191) ///< cannot allocate any more hlos mem
#define E_DPOOL_CAP_REACHED (AEE_EOFFSET + 0x192) ///< cannot allocate any more physpool mem
#define E_NO_MORE_FREE_SECTIONS (AEE_EOFFSET + 0x193) ///< No more free sections available to grow heap
/**
* @}
*/
#endif /* #ifndef AEESTDERR_H */

685
artifacts/package_sources/tinygrad/extra/dsp/include/HAP_power.h Executable file
View File

@@ -0,0 +1,685 @@
/*==============================================================================
@file
HAP_power.h
@brief
Header file of DSP power APIs.
Copyright (c) 2015,2019 Qualcomm Technologies, Inc.
All rights reserved. Qualcomm Proprietary and Confidential.
==============================================================================*/
#ifndef _HAP_POWER_H
#define _HAP_POWER_H
#include "AEEStdErr.h"
//#include <string.h>
//#include <stdlib.h>
#define boolean char
#define FALSE 0
#define TRUE 1
#define uint64 unsigned long long
#define uint32 unsigned int
#define NULL 0
#ifdef __cplusplus
extern "C" {
#endif
//Add a weak reference so shared objects do not throw link error
#pragma weak HAP_power_destroy_client
/**
* Possible error codes returned
*/
typedef enum {
HAP_POWER_ERR_UNKNOWN = -1,
HAP_POWER_ERR_INVALID_PARAM = -2,
HAP_POWER_ERR_UNSUPPORTED_API = -3
} HAP_power_error_codes;
/** Payload for HAP_power_set_mips_bw */
typedef struct {
boolean set_mips; /**< Set to TRUE to request MIPS */
unsigned int mipsPerThread; /**< mips requested per thread, to establish a minimal clock frequency per HW thread */
unsigned int mipsTotal; /**< Total mips requested, to establish total number of MIPS required across all HW threads */
boolean set_bus_bw; /**< Set to TRUE to request bus_bw */
uint64 bwBytePerSec; /**< Max bus BW requested (bytes per second) */
unsigned short busbwUsagePercentage; /**< Percentage of time during which bwBytesPerSec BW is required from the bus (0..100) */
boolean set_latency; /**< Set to TRUE to set latency */
int latency; /**< maximum hardware wakeup latency in microseconds. The higher the value,
* the deeper state of sleep that can be entered but the longer it may take
* to awaken. Only values > 0 are supported (1 microsecond is the smallest valid value) */
} HAP_power_mips_bw_payload;
/** @defgroup HAP_power_enums HAP POWER enums
* @{
*/
/** Clock frequency match type*/
typedef enum {
HAP_FREQ_AT_LEAST, /**< Matches at least the specified frequency. */
HAP_FREQ_AT_MOST, /**< Matches at most the specified frequency. */
HAP_FREQ_CLOSEST, /**< Closest match to the specified frequency. */
HAP_FREQ_EXACT, /**< Exact match with the specified frequency. */
HAP_FREQ_MAX_COUNT /**< Maximum count. */
} HAP_freq_match_type;
/**
* @} // HAP_power_enums
*/
/** Configuration for bus bandwidth */
typedef struct {
boolean set_bus_bw; /**< Set to TRUE to request bus_bw */
uint64 bwBytePerSec; /**< Max bus BW requested (bytes per second) */
unsigned short busbwUsagePercentage; /**< Percentage of time during which bwBytesPerSec BW is required from the bus (0..100) */
} HAP_power_bus_bw;
/**
* @brief Payload for vapps power request
* vapps core is used for Video post processing
*/
typedef struct {
boolean set_clk; /**< Set to TRUE to request clock frequency */
unsigned int clkFreqHz; /**< Clock frequency in Hz */
HAP_freq_match_type freqMatch; /**< Clock frequency match */
HAP_power_bus_bw dma_ext; /**< DMA external bus bandwidth */
HAP_power_bus_bw hcp_ext; /**< HCP external bus bandwidth */
HAP_power_bus_bw dma_int; /**< DMA internal bus bandwidth */
HAP_power_bus_bw hcp_int; /**< HCP internal bus bandwidth */
} HAP_power_vapss_payload;
/**
* @brief Payload for vapps_v2 power request
* Supported in targets which have split VAPPS core(DMA and HCP) form Hana onwards
*/
typedef struct {
boolean set_dma_clk; /**< Set to TRUE to reqeust DMA clock frequency */
boolean set_hcp_clk; /**< Set to TRUE to reqeust HCP clock frequency */
unsigned int dmaClkFreqHz; /**< DMA Clock frequency in Hz */
unsigned int hcpClkFreqHz; /**< HCP Clock frequency in Hz */
HAP_freq_match_type freqMatch; /**< Clock frequency match type */
HAP_power_bus_bw dma_ext; /**< DMA external bus bandwidth */
HAP_power_bus_bw hcp_ext; /**< HCP external bus bandwidth */
HAP_power_bus_bw dma_int; /**< DMA internal bus bandwidth */
HAP_power_bus_bw hcp_int; /**< HCP internal bus bandwidth */
} HAP_power_vapss_payload_v2;
/** Payload for HAP_power_set_HVX */
typedef struct {
boolean power_up; /**< Set to TRUE to turn on HVX, and FALSE to turn off. */
} HAP_power_hvx_payload;
/**
* Payload for HAP_power_set_HMX
* Supported from Lahaina onwards*/
typedef struct {
boolean power_up; /**< Set to TRUE to turn on HMX, and FALSE to turn off. */
} HAP_power_hmx_payload;
/** @defgroup HAP_power_enums HAP POWER enums
* @{
*/
/** Payload for HAP power client classes */
typedef enum {
HAP_POWER_UNKNOWN_CLIENT_CLASS = 0x00, /**< Unknown client class */
HAP_POWER_AUDIO_CLIENT_CLASS = 0x01, /**< Audio client class */
HAP_POWER_VOICE_CLIENT_CLASS = 0x02, /**< Voice client class */
HAP_POWER_COMPUTE_CLIENT_CLASS = 0x04, /**< Compute client class */
HAP_POWER_STREAMING_1HVX_CLIENT_CLASS = 0x08, /**< Camera streaming with 1 HVX client class */
HAP_POWER_STREAMING_2HVX_CLIENT_CLASS = 0x10, /**< Camera streaming with 2 HVX client class */
} HAP_power_app_type_payload;
/**
* @} // HAP_power_enums
*/
/** Payload for HAP_power_set_linelock */
typedef struct {
void* startAddress; /**< Start address of the memory region to be locked. */
uint32 size; /**< Size (bytes) of the memory region to be locked. Set size
* to 0 to unlock memory. */
uint32 throttleBlockSize; /**< Block size for throttling, in bytes;
* 0 for no throttling. The region to be locked will be divided into
* blocks of this size for throttling purposes.
* Use for locking larger cache blocks.
* Applicable only when enabling line locking.Only ONE throttled linelock call is supported at this time.
* You can linelock additional regions (without throttling) using HAP_power_set_linelock_nothrottle*/
uint32 throttlePauseUs; /**< Pause to be applied between locking each block, in microseconds. Applicable only when enabling line locking*/
} HAP_power_linelock_payload;
/** Payload for HAP_power_set_linelock_nothrottle */
typedef struct {
void* startAddress; /**< Start address of the memory region to be locked. */
uint32 size; /**< Size (bytes) of the memory region to be locked. Set size to 0
* to unlock memory */
} HAP_power_linelock_nothrottle_payload;
/** @defgroup HAP_power_enums HAP POWER enums
* @{
*/
/** Option for dcvs payload */
typedef enum {
HAP_DCVS_ADJUST_UP_DOWN = 0x1, /**< increase and decrease core/bus clock speed. */
HAP_DCVS_ADJUST_ONLY_UP = 0x2, /**< restricts DCVS from lowering the clock speed below the requested value . */
} HAP_power_dcvs_payload_option;
/**
* @} // HAP_power_enums
*/
/** Payload for HAP_power_set_DCVS */
typedef struct {
boolean dcvs_enable; /**< Set to TRUE to participate in DCVS, and FALSE otherwise. */
HAP_power_dcvs_payload_option dcvs_option; /**< Set to one of
* HAP_DCVS_ADJUST_UP_DOWN - Allows for DCVS to adjust up and down.
* HAP_DCVS_ADJUST_ONLY_UP - Allows for DCVS to adjust up only. */
} HAP_power_dcvs_payload;
/** @defgroup HAP_power_enums HAP POWER enums
* @{
*/
/** Voltage corners for HAP DCVS V2 interface */
typedef enum {
HAP_DCVS_VCORNER_DISABLE,
HAP_DCVS_VCORNER_SVS2,
HAP_DCVS_VCORNER_SVS,
HAP_DCVS_VCORNER_SVS_PLUS,
HAP_DCVS_VCORNER_NOM,
HAP_DCVS_VCORNER_NOM_PLUS,
HAP_DCVS_VCORNER_TURBO,
HAP_DCVS_VCORNER_TURBO_PLUS,
HAP_DCVS_VCORNER_MAX = 255,
} HAP_dcvs_voltage_corner_t;
/**
* @} // HAP_power_enums
*/
#define HAP_DCVS_VCORNER_SVSPLUS HAP_DCVS_VCORNER_SVS_PLUS
#define HAP_DCVS_VCORNER_NOMPLUS HAP_DCVS_VCORNER_NOM_PLUS
/** DCVS parameters for HAP_power_dcvs_v2_payload */
typedef struct {
HAP_dcvs_voltage_corner_t target_corner; /**< target voltage corner */
HAP_dcvs_voltage_corner_t min_corner; /**< minimum voltage corner */
HAP_dcvs_voltage_corner_t max_corner; /**< maximum voltage corner */
uint32 param1; /**< reserved */
uint32 param2; /**< reserved */
uint32 param3; /**< reserved */
} HAP_dcvs_params_t;
/** Core clock parameters for HAP_power_dcvs_v3_payload */
typedef struct {
HAP_dcvs_voltage_corner_t target_corner; /**< target voltage corner */
HAP_dcvs_voltage_corner_t min_corner; /**< minimum voltage corner */
HAP_dcvs_voltage_corner_t max_corner; /**< maximum voltage corner */
uint32 param1; /**< reserved */
uint32 param2; /**< reserved */
uint32 param3; /**< reserved */
} HAP_core_params_t;
/** Bus clock parameters for HAP_power_dcvs_v3_payload */
typedef struct {
HAP_dcvs_voltage_corner_t target_corner; /**< target voltage corner */
HAP_dcvs_voltage_corner_t min_corner; /**< minimum voltage corner */
HAP_dcvs_voltage_corner_t max_corner; /**< maximum voltage corner */
uint32 param1; /**< reserved */
uint32 param2; /**< reserved */
uint32 param3; /**< reserved */
} HAP_bus_params_t;
/** DCVS v3 parameters for HAP_power_dcvs_v3_payload */
typedef struct {
uint32 param1; /**< reserved */
uint32 param2; /**< reserved */
uint32 param3; /**< reserved */
uint32 param4; /**< reserved */
uint32 param5; /**< reserved */
uint32 param6; /**< reserved */
} HAP_dcvs_v3_params_t;
/** @defgroup HAP_power_enums HAP POWER enums
* @{
*/
/** option for dcvs_v2 payload */
typedef enum {
HAP_DCVS_V2_ADJUST_UP_DOWN = 0x1, /**< Allows for DCVS to adjust up and down. */
HAP_DCVS_V2_ADJUST_ONLY_UP = 0x2, /**< Allows for DCVS to adjust up only. */
HAP_DCVS_V2_POWER_SAVER_MODE = 0x4, /**< HAP_DCVS_POWER_SAVER_MODE - Higher thresholds for power efficiency. */
HAP_DCVS_V2_POWER_SAVER_AGGRESSIVE_MODE = 0x8, /**< HAP_DCVS_POWER_SAVER_AGGRESSIVE_MODE - Higher thresholds for power efficiency with faster ramp down. */
HAP_DCVS_V2_PERFORMANCE_MODE = 0x10, /**< HAP_DCVS_PERFORMANCE_MODE - Lower thresholds for maximum performance */
HAP_DCVS_V2_DUTY_CYCLE_MODE = 0x20, /**< HAP_DCVS_DUTY_CYCLE_MODE - only for HVX based clients.
* For streaming class clients:
* > detects periodicity based on HVX usage
* > lowers clocks in the no HVX activity region of each period.
* For compute class clients:
* > Lowers clocks on no HVX activity detects and brings clocks up on detecting HVX activity again.
* > Latency involved in bringing up the clock with be at max 1 to 2 ms. */
} HAP_power_dcvs_v2_payload_option;
/**
* @} // HAP_power_enums
*/
/** Payload for HAP_power_set_DCVS_v2 */
typedef struct {
boolean dcvs_enable; /**< Set to TRUE to participate in DCVS, and FALSE otherwise */
HAP_power_dcvs_v2_payload_option dcvs_option; /**< Set to one of HAP_power_dcvs_v2_payload_option */
boolean set_latency; /**< TRUE to set latency parameter, otherwise FALSE */
uint32 latency; /**< sleep latency */
boolean set_dcvs_params; /**< TRUE to set DCVS params, otherwise FALSE */
HAP_dcvs_params_t dcvs_params; /**< DCVS parameters */
} HAP_power_dcvs_v2_payload;
/** Payload for HAP_power_set_DCVS_v3 */
typedef struct {
boolean set_dcvs_enable; /**< TRUE to consider DCVS enable/disable and option parameters, otherwise FALSE */
boolean dcvs_enable; /**< Set to TRUE to participate in DCVS, and FALSE otherwise. */
HAP_power_dcvs_v2_payload_option dcvs_option; /**< Set to one of HAP_power_dcvs_v2_payload_option */
boolean set_latency; /**< TRUE to consider latency parameter, otherwise FALSE */
uint32 latency; /**< sleep latency */
boolean set_core_params; /**< TRUE to consider core clock params, otherwise FALSE */
HAP_core_params_t core_params; /**< Core clock parameters */
boolean set_bus_params; /**< TRUE to consider bus clock params, otherwise FALSE */
HAP_bus_params_t bus_params; /**< Bus clock parameters */
boolean set_dcvs_v3_params; /**< TRUE to consider DCVS v3 params, otherwise FALSE */
HAP_dcvs_v3_params_t dcvs_v3_params; /**< DCVS v3 parameters */
boolean set_sleep_disable; /**< TRUE to consider sleep disable/enable parameter, otherwise FALSE */
boolean sleep_disable; /**< TRUE to disable sleep/LPM modes, FALSE to enable */
} HAP_power_dcvs_v3_payload;
/** @defgroup HAP_power_enums HAP POWER enums
* @{
*/
/** Type for dcvs update request */
typedef enum {
HAP_POWER_UPDATE_DCVS = 1,
HAP_POWER_UPDATE_SLEEP_LATENCY,
HAP_POWER_UPDATE_DCVS_PARAMS,
} HAP_power_update_type_t;
/**
* @} // HAP_power_enums
*/
/** Payload for DCVS update */
typedef struct {
boolean dcvs_enable; /**< TRUE for DCVS enable and FALSE for DCVS disable */
HAP_power_dcvs_v2_payload_option dcvs_option; /**< Requested DCVS policy in case DCVS enable is TRUE */
} HAP_power_update_dcvs_t;
/** Payload for latency update */
typedef struct {
boolean set_latency; /**< TRUE if sleep latency request has to be considered */
unsigned int latency; /**< Sleep latency request in micro seconds */
} HAP_power_update_latency_t;
/** Payload for DCVS params update */
typedef struct {
boolean set_dcvs_params; /**< Flag to mark DCVS params structure validity, TRUE for valid DCVS
*params request and FALSE otherwise */
HAP_dcvs_params_t dcvs_params; /**< Intended DCVS params if set_dcvs_params is set to TRUE */
} HAP_power_update_dcvs_params_t;
/** Payload for HAP_power_set_DCVS_v2 */
typedef struct {
HAP_power_update_type_t update_param; /**< Type for which param to update */
union {
HAP_power_update_dcvs_t dcvs_payload;
HAP_power_update_latency_t latency_payload;
HAP_power_update_dcvs_params_t dcvs_params_payload;
}; /**< Update payload for DCVS, latency or DCVS params */
} HAP_power_dcvs_v2_update_payload;
/** Payload for HAP_power_set_streamer */
typedef struct {
boolean set_streamer0_clk; /**< Set streamer 0 clock */
boolean set_streamer1_clk; /**< Set streamer 1 clock */
unsigned int streamer0_clkFreqHz; /**< Streamer 0 clock frequency */
unsigned int streamer1_clkFreqHz; /**< Streamer 1 clock frequency */
HAP_freq_match_type freqMatch; /**< Clock frequency match */
uint32 param1; /**< Reserved for future streamer parameters */
uint32 param2; /**< Reserved for future streamer parameters */
uint32 param3; /**< Reserved for future streamer parameters */
} HAP_power_streamer_payload;
/** @defgroup HAP_power_enums HAP POWER enums
* @{
*/
/** Identifies the HAP power request type */
typedef enum {
HAP_power_set_mips_bw = 1, /**< Requests for MIPS. Provides
* fine-grained control to set MIPS values.
* Payload is set to HAP_power_payload */
HAP_power_set_HVX, /**< Requests to enable / disable HVX
* Payload is set to HAP_power_hvx_payload */
HAP_power_set_apptype, /**< Sets the app_type
* Payload is set to HAP_power_app_type_payload */
HAP_power_set_linelock, /**< Sets the throttled L2 cache line locking parameters.
* Only one throttled call is supported at this time. Additional
* un-throttled line-locks can be performed using HAP_power_set_linelock_nothrottle
* Payload is set to HAP_power_linelock_payload */
HAP_power_set_DCVS, /**< Requests to participate / stop participating in DCVS */
HAP_power_set_linelock_nothrottle, /**< Sets the L2 cache line locking parameters (non-throttled).
* Payload is set to HAP_power_linelock_nothrottle_payload */
HAP_power_set_DCVS_v2, /**< Requests to participate / stop participating in DCVS_v2 */
HAP_power_set_vapss, /**< Sets the VAPSS core clock and DDR/IPNOC bandwidth
* Payload is set to HAP_power_vapss_payload */
HAP_power_set_vapss_v2, /**< Sets the VAPSS core DMA/HCP clocks and DDR/IPNOC bandwidths
* Payload is set to HAP_power_vapss_payload_v2 */
HAP_power_set_dcvs_v2_update, /**< Updates DCVS params
* Payload is set to HAP_power_dcvs_v2_update_payload */
HAP_power_set_streamer, /**< Sets the streamer core clocks
* Payload is set to HAP_power_streamer_payload */
HAP_power_set_DCVS_v3, /**< Updates DCVS params
* Payload is set to HAP_power_dcvs_v3_payload */
HAP_power_set_HMX, /**< Requests to enable / disable HMX
* Payload is set to HAP_power_hmx_payload */
} HAP_Power_request_type;
/**
* @} // HAP_power_enums
*/
/** Data type to change power values on the DSP */
typedef struct {
HAP_Power_request_type type; /**< Identifies the request type */
union{
HAP_power_mips_bw_payload mips_bw; /**< Requests for performance level */
HAP_power_vapss_payload vapss; /**< Sets the VAPSS core clock and DDR/IPNOC bandwidth */
HAP_power_vapss_payload_v2 vapss_v2; /**< Sets the VAPSS core clock and DDR/IPNOC bandwidth */
HAP_power_streamer_payload streamer; /**< Sets the streamer core clocks */
HAP_power_hvx_payload hvx; /**< Requests to enable / disable HVX */
HAP_power_app_type_payload apptype; /**< Sets the app_type */
HAP_power_linelock_payload linelock; /**< Sets the throttled L2 cache linelock parameters. Only one
* throttled linelock is permitted at this time. Additional
* un-throttled linelocks can be performed using linelock_nothrottle */
HAP_power_dcvs_payload dcvs; /**< Updates DCVS params */
HAP_power_dcvs_v2_payload dcvs_v2; /**< Updates DCVS_v2 params */
HAP_power_dcvs_v2_update_payload dcvs_v2_update; /**< Updates DCVS_v2_update params */
HAP_power_linelock_nothrottle_payload linelock_nothrottle; /**< Sets the un-throttled L2 cache linelock parameters */
HAP_power_dcvs_v3_payload dcvs_v3; /**< Updates DCVS_v3 params */
HAP_power_hmx_payload hmx; /**< Requests to turn on / off HMX */
};
} HAP_power_request_t;
/** @defgroup HAP_power_functions HAP POWER functions
* @{
*/
/**
* Method to set power values from the DSP
* @param[in] context - To identify the power client
* @param[in] request - Request params.
* @retval 0 on success, AEE_EMMPMREGISTER on MMPM client register request failure, -1 on unknown error
*/
int HAP_power_set(void* context, HAP_power_request_t* request);
/**
* @} // HAP_power_functions
*/
/** @defgroup HAP_power_enums HAP POWER enums
* @{
*/
/** Identifies the HAP power response type */
typedef enum {
HAP_power_get_max_mips = 1, /**< Returns the max mips supported (max_mips) */
HAP_power_get_max_bus_bw, /**< Returns the max bus bandwidth supported (max_bus_bw) */
HAP_power_get_client_class, /**< Returns the client class (client_class) */
HAP_power_get_clk_Freq, /**< Returns the core clock frequency (clkFreqHz) */
HAP_power_get_aggregateAVSMpps, /**< Returns the aggregate Mpps used by audio and voice (clkFreqHz) */
HAP_power_get_dcvsEnabled, /**< Returns the dcvs status (enabled / disabled) */
HAP_power_get_vapss_core_clk_Freq, /**< Returns the VAPSS core clock frequency (clkFreqHz) */
HAP_power_get_dma_core_clk_Freq, /**< Returns the DMA core clock frequency (clkFreqHz) */
HAP_power_get_hcp_core_clk_Freq, /**< Returns the HCP core clock frequency (clkFreqHz) */
HAP_power_get_streamer0_core_clk_Freq, /**< Returns the streamer 0 core clock frequency (clkFreqHz) */
HAP_power_get_streamer1_core_clk_Freq, /**< Returns the streamer 1 core clock frequency (clkFreqHz) */
} HAP_Power_response_type;
/**
* @} // HAP_power_enums
*/
/** Data type to retrieve power values from the DSP */
typedef struct {
HAP_Power_response_type type; /**< Identifies the type to retrieve. */
union{
unsigned int max_mips; /**< Max mips supported */
uint64 max_bus_bw; /**< Max bus bw supported */
unsigned int client_class; /**< Current client class */
unsigned int clkFreqHz; /**< Current core CPU frequency */
unsigned int aggregateAVSMpps; /**< Aggregate AVS Mpps used by audio and voice */
boolean dcvsEnabled; /**< Indicates if dcvs is enabled / disabled. */
};
} HAP_power_response_t;
/** @defgroup HAP_power_functions HAP POWER functions
* @{
*/
/**
* Method to retrieve power values from the DSP
* @param[in] context - Ignored
* @param[out] response - Response.
*/
int HAP_power_get(void* context, HAP_power_response_t* response);
/**
* Method to initialize dcvs v3 structure in request param. It enables
* flags and resets params for all fields in dcvs v3. So, this
* can also be used to remove applied dcvs v3 params and restore
* defaults.
* @param[in] request - Pointer to request params.
*/
/*static inline void HAP_power_set_dcvs_v3_init(HAP_power_request_t* request) {
memset(request, 0, sizeof(HAP_power_request_t) );
request->type = HAP_power_set_DCVS_v3;
request->dcvs_v3.set_dcvs_enable = TRUE;
request->dcvs_v3.dcvs_enable = TRUE;
request->dcvs_v3.dcvs_option = HAP_DCVS_V2_POWER_SAVER_MODE;
request->dcvs_v3.set_latency = TRUE;
request->dcvs_v3.latency = 65535;
request->dcvs_v3.set_core_params = TRUE;
request->dcvs_v3.set_bus_params = TRUE;
request->dcvs_v3.set_dcvs_v3_params = TRUE;
request->dcvs_v3.set_sleep_disable = TRUE;
return;
}*/
/**
* Method to enable/disable dcvs and set particular dcvs policy.
* @param[in] context - User context.
* @param[in] dcvs_enable - TRUE to enable dcvs, FALSE to disable dcvs.
* @param[in] dcvs_option - To set particular dcvs policy. In case of dcvs disable
* request, this param will be ignored.
* @returns - 0 on success
*/
/*static inline int HAP_power_set_dcvs_option(void* context, boolean dcvs_enable,
HAP_power_dcvs_v2_payload_option dcvs_option) {
HAP_power_request_t request;
memset(&request, 0, sizeof(HAP_power_request_t) );
request.type = HAP_power_set_DCVS_v3;
request.dcvs_v3.set_dcvs_enable = TRUE;
request.dcvs_v3.dcvs_enable = dcvs_enable;
if(dcvs_enable)
request.dcvs_v3.dcvs_option = dcvs_option;
return HAP_power_set(context, &request);
}*/
/**
* Method to set/reset sleep latency.
* @param[in] context - User context.
* @param[in] latency - Sleep latency value in microseconds, should be > 1.
* Use 65535 max value to reset it to default.
* @returns - 0 on success
*/
/*static inline int HAP_power_set_sleep_latency(void* context, uint32 latency) {
HAP_power_request_t request;
memset(&request, 0, sizeof(HAP_power_request_t) );
request.type = HAP_power_set_DCVS_v3;
request.dcvs_v3.set_latency = TRUE;
request.dcvs_v3.latency = latency;
return HAP_power_set(context, &request);
}*/
/**
* Method to set/reset DSP core clock voltage corners.
* @param[in] context - User context.
* @param[in] target_corner - Target voltage corner.
* @param[in] min_corner - Minimum voltage corner.
* @param[in] max_corner - Maximum voltage corner.
* @returns - 0 on success
*/
/*static inline int HAP_power_set_core_corner(void* context, uint32 target_corner,
uint32 min_corner, uint32 max_corner) {
HAP_power_request_t request;
memset(&request, 0, sizeof(HAP_power_request_t) );
request.type = HAP_power_set_DCVS_v3;
request.dcvs_v3.set_core_params = TRUE;
request.dcvs_v3.core_params.min_corner = (HAP_dcvs_voltage_corner_t) (min_corner);
request.dcvs_v3.core_params.max_corner = (HAP_dcvs_voltage_corner_t) (max_corner);
request.dcvs_v3.core_params.target_corner = (HAP_dcvs_voltage_corner_t) (target_corner);
return HAP_power_set(context, &request);
}*/
/**
* Method to set/reset bus clock voltage corners.
* @param[in] context - User context.
* @param[in] target_corner - Target voltage corner.
* @param[in] min_corner - Minimum voltage corner.
* @param[in] max_corner - Maximum voltage corner.
* @returns - 0 on success
*/
/*static inline int HAP_power_set_bus_corner(void* context, uint32 target_corner,
uint32 min_corner, uint32 max_corner) {
HAP_power_request_t request;
memset(&request, 0, sizeof(HAP_power_request_t) );
request.type = HAP_power_set_DCVS_v3;
request.dcvs_v3.set_bus_params = TRUE;
request.dcvs_v3.bus_params.min_corner = (HAP_dcvs_voltage_corner_t) (min_corner);
request.dcvs_v3.bus_params.max_corner = (HAP_dcvs_voltage_corner_t) (max_corner);
request.dcvs_v3.bus_params.target_corner = (HAP_dcvs_voltage_corner_t) (target_corner);
return HAP_power_set(context, &request);
}*/
/**
* Method to disable/enable all low power modes.
* @param[in] context - User context.
* @param[in] sleep_disable - TRUE to disable all low power modes.
* FALSE to re-enable all low power modes.
* @returns - 0 on success
*/
/*static inline int HAP_power_set_sleep_mode(void* context, boolean sleep_disable) {
HAP_power_request_t request;
memset(&request, 0, sizeof(HAP_power_request_t) );
request.type = HAP_power_set_DCVS_v3;
request.dcvs_v3.set_sleep_disable = TRUE;
request.dcvs_v3.sleep_disable = sleep_disable;
return HAP_power_set(context, &request);
}*/
/**
* This API is deprecated and might generate undesired results.
* Please use the HAP_power_get() and HAP_power_set() APIs instead.
* Requests a performance level by percentage for clock speed
* and bus speed. Passing 0 for any parameter results in no
* request being issued for that particular attribute.
* @param[in] clock - percentage of target's maximum clock speed
* @param[in] bus - percentage of target's maximum bus speed
* @param[in] latency - maximum hardware wake up latency in microseconds. The
* higher the value the deeper state of sleep
* that can be entered but the longer it may
* take to awaken.
* @retval 0 on success
* @par Comments : Performance metrics vary from target to target so the
* intent of this API is to allow callers to set a relative
* performance level to achieve the desired balance between
* performance and power saving.
*/
int HAP_power_request(int clock, int bus, int latency);
/**
* This API is deprecated and might generate undesired results.
* Please use the HAP_power_get() and HAP_power_set() APIs instead.
* Requests a performance level by absolute values. Passing 0
* for any parameter results in no request being issued for that
* particular attribute.
* @param[in] clock - speed in MHz
* @param[in] bus - bus speed in MHz
* @param[in] latency - maximum hardware wakeup latency in microseconds. The
* higher the value the deeper state of
* sleep that can be entered but the
* longer it may take to awaken.
* @retval 0 on success
* @par Comments : This API allows callers who are aware of their target
* specific capabilities to set them explicitly.
*/
int HAP_power_request_abs(int clock, int bus, int latency);
/**
* This API is deprecated and might generate undesired results.
* Please use the HAP_power_get() and HAP_power_set() APIs instead.
* queries the target for its clock and bus speed capabilities
* @param[out] clock_max - maximum clock speed supported in MHz
* @param[out] bus_max - maximum bus speed supported in MHz
* @retval 0 on success
*/
int HAP_power_get_max_speed(int* clock_max, int* bus_max);
/**
* This API is deprecated and might generate undesired results.
* Please use the HAP_power_get() and HAP_power_set() APIs instead.
* Upvote for HVX power
* @retval 0 on success
*/
int HVX_power_request(void);
/**
* This API is deprecated and might generate undesired results.
* Please use the HAP_power_get() and HAP_power_set() APIs instead.
* Downvote for HVX power
* @retval 0 on success
*/
int HVX_power_release(void);
/**
* Method to destroy clients created through HAP_power_set
* @param[in] context - To uniquely identify the client
* @retval 0 on success, AEE_ENOSUCHCLIENT on Invalid context, -1 on unknown error
* @brief DO NOT call this API directly, use HAP_power_destroy instead.
*/
int HAP_power_destroy_client(void *context);
/**
* @param[in] client - To uniquely identify the client context.
* @retval 0 on success, AEE_EUNSUPPORTEDAPI if the API is not supported on the DSP image, AEE_ENOSUCHCLIENT on Invalid context, -1 on unknown error
* @brief Method to destroy clients created through HAP_power_set, wrapper to HAP_power_destroy_client API
*/
static inline int HAP_power_destroy(void *client){
if(0 != HAP_power_destroy_client)
return HAP_power_destroy_client(client);
return AEE_EUNSUPPORTEDAPI;
}
/**
* Method to create user client context
* @retval context for client
*/
//static inline void* HAP_utils_create_context(void) {
/*
* Allocate 1 byte of memory for a unique context identifier
* Clients can also allocate memory and use it as unique context identifier
*/
// return malloc(1);
//}
/**
* Method to destroy user client context
* @param context of client
*/
/*static inline void HAP_utils_destroy_context(void* context) {
free(context);
}*/
/**
* @} // HAP_power_functions
*/
#ifdef __cplusplus
}
#endif
#endif //_HAP_POWER_H

319
artifacts/package_sources/tinygrad/extra/dsp/include/adsprpc_shared.h Normal file
View File

@@ -0,0 +1,319 @@
/*
* Copyright (c) 2012-2018, The Linux Foundation. All rights reserved.
*
* This program is free software; you can redistribute it and/or modify
* it under the terms of the GNU General Public License version 2 and
* only version 2 as published by the Free Software Foundation.
*
* This program is distributed in the hope that it will be useful,
* but WITHOUT ANY WARRANTY; without even the implied warranty of
* MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
* GNU General Public License for more details.
*
*/
#ifndef ADSPRPC_SHARED_H
#define ADSPRPC_SHARED_H
#include <stdint.h>
#include <stddef.h>
#include <sys/types.h>
#include <linux/types.h>
#define FASTRPC_IOCTL_INVOKE _IOWR('R', 1, struct fastrpc_ioctl_invoke)
#define FASTRPC_IOCTL_MMAP _IOWR('R', 2, struct fastrpc_ioctl_mmap)
#define FASTRPC_IOCTL_MUNMAP _IOWR('R', 3, struct fastrpc_ioctl_munmap)
#define FASTRPC_IOCTL_MMAP_64 _IOWR('R', 14, struct fastrpc_ioctl_mmap_64)
#define FASTRPC_IOCTL_MUNMAP_64 _IOWR('R', 15, struct fastrpc_ioctl_munmap_64)
#define FASTRPC_IOCTL_INVOKE_FD _IOWR('R', 4, struct fastrpc_ioctl_invoke_fd)
#define FASTRPC_IOCTL_SETMODE _IOWR('R', 5, uint32_t)
#define FASTRPC_IOCTL_INIT _IOWR('R', 6, struct fastrpc_ioctl_init)
#define FASTRPC_IOCTL_INVOKE_ATTRS \
_IOWR('R', 7, struct fastrpc_ioctl_invoke_attrs)
#define FASTRPC_IOCTL_GETINFO _IOWR('R', 8, uint32_t)
#define FASTRPC_IOCTL_GETPERF _IOWR('R', 9, struct fastrpc_ioctl_perf)
#define FASTRPC_IOCTL_INIT_ATTRS _IOWR('R', 10, struct fastrpc_ioctl_init_attrs)
#define FASTRPC_IOCTL_INVOKE_CRC _IOWR('R', 11, struct fastrpc_ioctl_invoke_crc)
#define FASTRPC_IOCTL_CONTROL _IOWR('R', 12, struct fastrpc_ioctl_control)
#define FASTRPC_IOCTL_MUNMAP_FD _IOWR('R', 13, struct fastrpc_ioctl_munmap_fd)
#define FASTRPC_GLINK_GUID "fastrpcglink-apps-dsp"
#define FASTRPC_SMD_GUID "fastrpcsmd-apps-dsp"
#define DEVICE_NAME "adsprpc-smd"
/* Set for buffers that have no virtual mapping in userspace */
#define FASTRPC_ATTR_NOVA 0x1
/* Set for buffers that are NOT dma coherent */
#define FASTRPC_ATTR_NON_COHERENT 0x2
/* Set for buffers that are dma coherent */
#define FASTRPC_ATTR_COHERENT 0x4
/* Fastrpc attribute for keeping the map persistent */
#define FASTRPC_ATTR_KEEP_MAP 0x8
/* Fastrpc attribute for no map */
#define FASTRPC_ATTR_NOMAP (16)
/* Driver should operate in parallel with the co-processor */
#define FASTRPC_MODE_PARALLEL 0
/* Driver should operate in serial mode with the co-processor */
#define FASTRPC_MODE_SERIAL 1
/* Driver should operate in profile mode with the co-processor */
#define FASTRPC_MODE_PROFILE 2
/* Set FastRPC session ID to 1 */
#define FASTRPC_MODE_SESSION 4
/* INIT a new process or attach to guestos */
#define FASTRPC_INIT_ATTACH 0
#define FASTRPC_INIT_CREATE 1
#define FASTRPC_INIT_CREATE_STATIC 2
#define FASTRPC_INIT_ATTACH_SENSORS 3
/* Retrives number of input buffers from the scalars parameter */
#define REMOTE_SCALARS_INBUFS(sc) (((sc) >> 16) & 0x0ff)
/* Retrives number of output buffers from the scalars parameter */
#define REMOTE_SCALARS_OUTBUFS(sc) (((sc) >> 8) & 0x0ff)
/* Retrives number of input handles from the scalars parameter */
#define REMOTE_SCALARS_INHANDLES(sc) (((sc) >> 4) & 0x0f)
/* Retrives number of output handles from the scalars parameter */
#define REMOTE_SCALARS_OUTHANDLES(sc) ((sc) & 0x0f)
#define REMOTE_SCALARS_LENGTH(sc) (REMOTE_SCALARS_INBUFS(sc) +\
REMOTE_SCALARS_OUTBUFS(sc) +\
REMOTE_SCALARS_INHANDLES(sc) +\
REMOTE_SCALARS_OUTHANDLES(sc))
#define REMOTE_SCALARS_MAKEX(attr, method, in, out, oin, oout) \
((((uint32_t) (attr) & 0x7) << 29) | \
(((uint32_t) (method) & 0x1f) << 24) | \
(((uint32_t) (in) & 0xff) << 16) | \
(((uint32_t) (out) & 0xff) << 8) | \
(((uint32_t) (oin) & 0x0f) << 4) | \
((uint32_t) (oout) & 0x0f))
#define REMOTE_SCALARS_MAKE(method, in, out) \
REMOTE_SCALARS_MAKEX(0, method, in, out, 0, 0)
#ifndef VERIFY_PRINT_ERROR
#define VERIFY_EPRINTF(format, args) (void)0
#endif
#ifndef VERIFY_PRINT_INFO
#define VERIFY_IPRINTF(args) (void)0
#endif
#ifndef VERIFY
#define __STR__(x) #x ":"
#define __TOSTR__(x) __STR__(x)
#define __FILE_LINE__ __FILE__ ":" __TOSTR__(__LINE__)
#define VERIFY(err, val) \
do {\
VERIFY_IPRINTF(__FILE_LINE__"info: calling: " #val "\n");\
if ((val) == 0) {\
(err) = (err) == 0 ? -1 : (err);\
VERIFY_EPRINTF(__FILE_LINE__"error: %d: " #val "\n", (err));\
} else {\
VERIFY_IPRINTF(__FILE_LINE__"info: passed: " #val "\n");\
} \
} while (0)
#endif
#define remote_arg64_t union remote_arg64
struct remote_buf64 {
uint64_t pv;
uint64_t len;
};
struct remote_dma_handle64 {
int fd;
uint32_t offset;
uint32_t len;
};
union remote_arg64 {
struct remote_buf64 buf;
struct remote_dma_handle64 dma;
uint32_t h;
};
#define remote_arg_t union remote_arg
struct remote_buf {
void *pv; /* buffer pointer */
size_t len; /* length of buffer */
};
struct remote_dma_handle {
int fd;
uint32_t offset;
};
union remote_arg {
struct remote_buf buf; /* buffer info */
struct remote_dma_handle dma;
uint32_t h; /* remote handle */
};
struct fastrpc_ioctl_invoke {
uint32_t handle; /* remote handle */
uint32_t sc; /* scalars describing the data */
remote_arg_t *pra; /* remote arguments list */
};
struct fastrpc_ioctl_invoke_fd {
struct fastrpc_ioctl_invoke inv;
int *fds; /* fd list */
};
struct fastrpc_ioctl_invoke_attrs {
struct fastrpc_ioctl_invoke inv;
int *fds; /* fd list */
unsigned int *attrs; /* attribute list */
};
struct fastrpc_ioctl_invoke_crc {
struct fastrpc_ioctl_invoke inv;
int *fds; /* fd list */
unsigned int *attrs; /* attribute list */
unsigned int *crc;
};
struct fastrpc_ioctl_init {
uint32_t flags; /* one of FASTRPC_INIT_* macros */
uintptr_t file; /* pointer to elf file */
uint32_t filelen; /* elf file length */
int32_t filefd; /* ION fd for the file */
uintptr_t mem; /* mem for the PD */
uint32_t memlen; /* mem length */
int32_t memfd; /* ION fd for the mem */
};
struct fastrpc_ioctl_init_attrs {
struct fastrpc_ioctl_init init;
int attrs;
unsigned int siglen;
};
struct fastrpc_ioctl_munmap {
uintptr_t vaddrout; /* address to unmap */
size_t size; /* size */
};
struct fastrpc_ioctl_munmap_64 {
uint64_t vaddrout; /* address to unmap */
size_t size; /* size */
};
struct fastrpc_ioctl_mmap {
int fd; /* ion fd */
uint32_t flags; /* flags for dsp to map with */
uintptr_t vaddrin; /* optional virtual address */
size_t size; /* size */
uintptr_t vaddrout; /* dsps virtual address */
};
struct fastrpc_ioctl_mmap_64 {
int fd; /* ion fd */
uint32_t flags; /* flags for dsp to map with */
uint64_t vaddrin; /* optional virtual address */
size_t size; /* size */
uint64_t vaddrout; /* dsps virtual address */
};
struct fastrpc_ioctl_munmap_fd {
int fd; /* fd */
uint32_t flags; /* control flags */
uintptr_t va; /* va */
ssize_t len; /* length */
};
struct fastrpc_ioctl_perf { /* kernel performance data */
uintptr_t data;
uint32_t numkeys;
uintptr_t keys;
};
#define FASTRPC_CONTROL_LATENCY (1)
struct fastrpc_ctrl_latency {
uint32_t enable; /* latency control enable */
uint32_t level; /* level of control */
};
#define FASTRPC_CONTROL_SMMU (2)
struct fastrpc_ctrl_smmu {
uint32_t sharedcb;
};
#define FASTRPC_CONTROL_KALLOC (3)
struct fastrpc_ctrl_kalloc {
uint32_t kalloc_support; /* Remote memory allocation from kernel */
};
struct fastrpc_ioctl_control {
uint32_t req;
union {
struct fastrpc_ctrl_latency lp;
struct fastrpc_ctrl_smmu smmu;
struct fastrpc_ctrl_kalloc kalloc;
};
};
struct smq_null_invoke {
uint64_t ctx; /* invoke caller context */
uint32_t handle; /* handle to invoke */
uint32_t sc; /* scalars structure describing the data */
};
struct smq_phy_page {
uint64_t addr; /* physical address */
uint64_t size; /* size of contiguous region */
};
struct smq_invoke_buf {
int num; /* number of contiguous regions */
int pgidx; /* index to start of contiguous region */
};
struct smq_invoke {
struct smq_null_invoke header;
struct smq_phy_page page; /* remote arg and list of pages address */
};
struct smq_msg {
uint32_t pid; /* process group id */
uint32_t tid; /* thread id */
struct smq_invoke invoke;
};
struct smq_invoke_rsp {
uint64_t ctx; /* invoke caller context */
int retval; /* invoke return value */
};
static inline struct smq_invoke_buf *smq_invoke_buf_start(remote_arg64_t *pra,
uint32_t sc)
{
unsigned int len = REMOTE_SCALARS_LENGTH(sc);
return (struct smq_invoke_buf *)(&pra[len]);
}
static inline struct smq_phy_page *smq_phy_page_start(uint32_t sc,
struct smq_invoke_buf *buf)
{
unsigned int nTotal = REMOTE_SCALARS_LENGTH(sc);
return (struct smq_phy_page *)(&buf[nTotal]);
}
#endif

208
artifacts/package_sources/tinygrad/extra/dsp/include/apps_std.h Normal file
View File

@@ -0,0 +1,208 @@
/**
* Copyright (c) 2019, The Linux Foundation. All rights reserved.
*
* Redistribution and use in source and binary forms, with or without
* modification, are permitted provided that the following conditions are
* met:
* * Redistributions of source code must retain the above copyright
* notice, this list of conditions and the following disclaimer.
* * Redistributions in binary form must reproduce the above
* copyright notice, this list of conditions and the following
* disclaimer in the documentation and/or other materials provided
* with the distribution.
* * Neither the name of The Linux Foundation nor the names of its
* contributors may be used to endorse or promote products derived
* from this software without specific prior written permission.
*
* THIS SOFTWARE IS PROVIDED "AS IS" AND ANY EXPRESS OR IMPLIED
* WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF
* MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NON-INFRINGEMENT
* ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS
* BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
* CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
* SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR
* BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY,
* WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE
* OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN
* IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
*/
#ifndef _APPS_STD_H
#define _APPS_STD_H
#include "AEEStdDef.h"
#ifndef __QAIC_HEADER
#define __QAIC_HEADER(ff) ff
#endif //__QAIC_HEADER
#ifndef __QAIC_HEADER_EXPORT
#define __QAIC_HEADER_EXPORT
#endif // __QAIC_HEADER_EXPORT
#ifndef __QAIC_HEADER_ATTRIBUTE
#define __QAIC_HEADER_ATTRIBUTE
#endif // __QAIC_HEADER_ATTRIBUTE
#ifndef __QAIC_IMPL
#define __QAIC_IMPL(ff) ff
#endif //__QAIC_IMPL
#ifndef __QAIC_IMPL_EXPORT
#define __QAIC_IMPL_EXPORT
#endif // __QAIC_IMPL_EXPORT
#ifndef __QAIC_IMPL_ATTRIBUTE
#define __QAIC_IMPL_ATTRIBUTE
#endif // __QAIC_IMPL_ATTRIBUTE
#ifdef __cplusplus
extern "C" {
#endif
#if !defined(__QAIC_STRING1_OBJECT_DEFINED__) && !defined(__STRING1_OBJECT__)
#define __QAIC_STRING1_OBJECT_DEFINED__
#define __STRING1_OBJECT__
typedef struct _cstring1_s {
char* data;
int dataLen;
} _cstring1_t;
#endif /* __QAIC_STRING1_OBJECT_DEFINED__ */
/**
* standard library functions remoted from the apps to the dsp
*/
typedef int apps_std_FILE;
enum apps_std_SEEK {
APPS_STD_SEEK_SET,
APPS_STD_SEEK_CUR,
APPS_STD_SEEK_END,
_32BIT_PLACEHOLDER_apps_std_SEEK = 0x7fffffff
};
typedef enum apps_std_SEEK apps_std_SEEK;
typedef struct apps_std_DIR apps_std_DIR;
struct apps_std_DIR {
uint64 handle;
};
typedef struct apps_std_DIRENT apps_std_DIRENT;
struct apps_std_DIRENT {
int ino;
char name[255];
};
typedef struct apps_std_STAT apps_std_STAT;
struct apps_std_STAT {
uint64 tsz;
uint64 dev;
uint64 ino;
uint32 mode;
uint32 nlink;
uint64 rdev;
uint64 size;
int64 atime;
int64 atimensec;
int64 mtime;
int64 mtimensec;
int64 ctime;
int64 ctimensec;
};
/**
* @retval, if operation fails errno is returned
*/
__QAIC_HEADER_EXPORT int __QAIC_HEADER(apps_std_fopen)(const char* name, const char* mode, apps_std_FILE* psout) __QAIC_HEADER_ATTRIBUTE;
__QAIC_HEADER_EXPORT int __QAIC_HEADER(apps_std_freopen)(apps_std_FILE sin, const char* name, const char* mode, apps_std_FILE* psout) __QAIC_HEADER_ATTRIBUTE;
__QAIC_HEADER_EXPORT int __QAIC_HEADER(apps_std_fflush)(apps_std_FILE sin) __QAIC_HEADER_ATTRIBUTE;
__QAIC_HEADER_EXPORT int __QAIC_HEADER(apps_std_fclose)(apps_std_FILE sin) __QAIC_HEADER_ATTRIBUTE;
/**
* @param, bEOF, if read or write bytes <= bufLen bytes then feof() is called
* and the result is returned in bEOF, otherwise bEOF is set to 0.
* @retval, if read or write return 0 for non zero length buffers, ferror is checked
* and a non zero value is returned in case of error with no rout parameters
*/
__QAIC_HEADER_EXPORT int __QAIC_HEADER(apps_std_fread)(apps_std_FILE sin, byte* buf, int bufLen, int* bytesRead, int* bEOF) __QAIC_HEADER_ATTRIBUTE;
__QAIC_HEADER_EXPORT int __QAIC_HEADER(apps_std_fwrite)(apps_std_FILE sin, const byte* buf, int bufLen, int* bytesWritten, int* bEOF) __QAIC_HEADER_ATTRIBUTE;
/**
* @param, pos, this buffer is filled up to MIN(posLen, sizeof(fpos_t))
* @param, posLenReq, returns sizeof(fpos_t)
* @retval, if operation fails errno is returned
*/
__QAIC_HEADER_EXPORT int __QAIC_HEADER(apps_std_fgetpos)(apps_std_FILE sin, byte* pos, int posLen, int* posLenReq) __QAIC_HEADER_ATTRIBUTE;
/**
* @param, if size of pos doesn't match the system size an error is returned.
* fgetpos can be used to query the size of fpos_t
* @retval, if operation fails errno is returned
*/
__QAIC_HEADER_EXPORT int __QAIC_HEADER(apps_std_fsetpos)(apps_std_FILE sin, const byte* pos, int posLen) __QAIC_HEADER_ATTRIBUTE;
/**
* @retval, if operation fails errno is returned
*/
__QAIC_HEADER_EXPORT int __QAIC_HEADER(apps_std_ftell)(apps_std_FILE sin, int* pos) __QAIC_HEADER_ATTRIBUTE;
__QAIC_HEADER_EXPORT int __QAIC_HEADER(apps_std_fseek)(apps_std_FILE sin, int offset, apps_std_SEEK whence) __QAIC_HEADER_ATTRIBUTE;
__QAIC_HEADER_EXPORT int __QAIC_HEADER(apps_std_flen)(apps_std_FILE sin, uint64* len) __QAIC_HEADER_ATTRIBUTE;
/**
* @retval, only fails if transport fails
*/
__QAIC_HEADER_EXPORT int __QAIC_HEADER(apps_std_rewind)(apps_std_FILE sin) __QAIC_HEADER_ATTRIBUTE;
__QAIC_HEADER_EXPORT int __QAIC_HEADER(apps_std_feof)(apps_std_FILE sin, int* bEOF) __QAIC_HEADER_ATTRIBUTE;
__QAIC_HEADER_EXPORT int __QAIC_HEADER(apps_std_ferror)(apps_std_FILE sin, int* err) __QAIC_HEADER_ATTRIBUTE;
__QAIC_HEADER_EXPORT int __QAIC_HEADER(apps_std_clearerr)(apps_std_FILE sin) __QAIC_HEADER_ATTRIBUTE;
__QAIC_HEADER_EXPORT int __QAIC_HEADER(apps_std_print_string)(const char* str) __QAIC_HEADER_ATTRIBUTE;
/**
* @param val, must contain space for NULL
* @param valLenReq, length required with NULL
* @retval, if fails errno is returned
*/
__QAIC_HEADER_EXPORT int __QAIC_HEADER(apps_std_getenv)(const char* name, char* val, int valLen, int* valLenReq) __QAIC_HEADER_ATTRIBUTE;
/**
* @retval, if fails errno is returned
*/
__QAIC_HEADER_EXPORT int __QAIC_HEADER(apps_std_setenv)(const char* name, const char* val, int override) __QAIC_HEADER_ATTRIBUTE;
__QAIC_HEADER_EXPORT int __QAIC_HEADER(apps_std_unsetenv)(const char* name) __QAIC_HEADER_ATTRIBUTE;
/**
* This function will try to open a file given directories in envvarname separated by
* delim.
* so given environment variable FOO_PATH=/foo;/bar
* fopen_wth_env("FOO_PATH", ";", "path/to/file", "rw", &out);
* will try to open /foo/path/to/file, /bar/path/to/file
* if the variable is unset, it will open the file directly
*
* @param envvarname, name of the environment variable containing the path
* @param delim, delimiator string, such as ";"
* @param name, name of the file
* @param mode, mode
* @param psout, output handle
* @retval, 0 on success errno or -1 on failure
*/
__QAIC_HEADER_EXPORT int __QAIC_HEADER(apps_std_fopen_with_env)(const char* envvarname, const char* delim, const char* name, const char* mode, apps_std_FILE* psout) __QAIC_HEADER_ATTRIBUTE;
__QAIC_HEADER_EXPORT int __QAIC_HEADER(apps_std_fgets)(apps_std_FILE sin, byte* buf, int bufLen, int* bEOF) __QAIC_HEADER_ATTRIBUTE;
/**
* This method will return the paths that are searched when looking for a file.
* The paths are defined by the environment variable (separated by delimiters)
* that is passed to the method.
*
* @param envvarname, name of the environment variable containing the path
* @param delim, delimiator string, such as ";"
* @param name, name of the file
* @param paths, Search paths
* @param numPaths, Actual number of paths found
* @param maxPathLen, The max path length
* @retval, 0 on success errno or -1 on failure
*
*/
__QAIC_HEADER_EXPORT int __QAIC_HEADER(apps_std_get_search_paths_with_env)(const char* envvarname, const char* delim, _cstring1_t* paths, int pathsLen, uint32* numPaths, uint16* maxPathLen) __QAIC_HEADER_ATTRIBUTE;
__QAIC_HEADER_EXPORT int __QAIC_HEADER(apps_std_fileExists)(const char* path, boolean* exists) __QAIC_HEADER_ATTRIBUTE;
__QAIC_HEADER_EXPORT int __QAIC_HEADER(apps_std_fsync)(apps_std_FILE sin) __QAIC_HEADER_ATTRIBUTE;
__QAIC_HEADER_EXPORT int __QAIC_HEADER(apps_std_fremove)(const char* name) __QAIC_HEADER_ATTRIBUTE;
/**
* This function decrypts the file using the provided open file descriptor, closes the
* original descriptor and return a new file descriptor.
* @retval, if operation fails errno is returned
*/
__QAIC_HEADER_EXPORT int __QAIC_HEADER(apps_std_fdopen_decrypt)(apps_std_FILE sin, apps_std_FILE* psout) __QAIC_HEADER_ATTRIBUTE;
__QAIC_HEADER_EXPORT int __QAIC_HEADER(apps_std_opendir)(const char* name, apps_std_DIR* dir) __QAIC_HEADER_ATTRIBUTE;
__QAIC_HEADER_EXPORT int __QAIC_HEADER(apps_std_closedir)(const apps_std_DIR* dir) __QAIC_HEADER_ATTRIBUTE;
__QAIC_HEADER_EXPORT int __QAIC_HEADER(apps_std_readdir)(const apps_std_DIR* dir, apps_std_DIRENT* dirent, int* bEOF) __QAIC_HEADER_ATTRIBUTE;
__QAIC_HEADER_EXPORT int __QAIC_HEADER(apps_std_mkdir)(const char* name, int mode) __QAIC_HEADER_ATTRIBUTE;
__QAIC_HEADER_EXPORT int __QAIC_HEADER(apps_std_rmdir)(const char* name) __QAIC_HEADER_ATTRIBUTE;
__QAIC_HEADER_EXPORT int __QAIC_HEADER(apps_std_stat)(const char* name, apps_std_STAT* stat) __QAIC_HEADER_ATTRIBUTE;
__QAIC_HEADER_EXPORT int __QAIC_HEADER(apps_std_ftrunc)(apps_std_FILE sin, int offset) __QAIC_HEADER_ATTRIBUTE;
__QAIC_HEADER_EXPORT int __QAIC_HEADER(apps_std_frename)(const char* oldname, const char* newname) __QAIC_HEADER_ATTRIBUTE;
#ifdef __cplusplus
}
#endif
#endif //_APPS_STD_H

204
artifacts/package_sources/tinygrad/extra/dsp/include/ion.h Normal file
View File

@@ -0,0 +1,204 @@
/*
* drivers/staging/android/uapi/ion.h
*
* Copyright (C) 2011 Google, Inc.
*
* This software is licensed under the terms of the GNU General Public
* License version 2, as published by the Free Software Foundation, and
* may be copied, distributed, and modified under those terms.
*
* This program is distributed in the hope that it will be useful,
* but WITHOUT ANY WARRANTY; without even the implied warranty of
* MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
* GNU General Public License for more details.
*
*/
#ifndef _UAPI_LINUX_ION_H
#define _UAPI_LINUX_ION_H
#include <stddef.h>
#include <linux/ioctl.h>
#include <linux/types.h>
typedef int ion_user_handle_t;
/**
* enum ion_heap_types - list of all possible types of heaps
* @ION_HEAP_TYPE_SYSTEM: memory allocated via vmalloc
* @ION_HEAP_TYPE_SYSTEM_CONTIG: memory allocated via kmalloc
* @ION_HEAP_TYPE_CARVEOUT: memory allocated from a prereserved
* carveout heap, allocations are physically
* contiguous
* @ION_HEAP_TYPE_DMA: memory allocated via DMA API
* @ION_NUM_HEAPS: helper for iterating over heaps, a bit mask
* is used to identify the heaps, so only 32
* total heap types are supported
*/
enum ion_heap_type {
ION_HEAP_TYPE_SYSTEM,
ION_HEAP_TYPE_SYSTEM_CONTIG,
ION_HEAP_TYPE_CARVEOUT,
ION_HEAP_TYPE_CHUNK,
ION_HEAP_TYPE_DMA,
ION_HEAP_TYPE_CUSTOM, /*
* must be last so device specific heaps always
* are at the end of this enum
*/
ION_NUM_HEAPS = 16,
};
#define ION_HEAP_SYSTEM_MASK ((1 << ION_HEAP_TYPE_SYSTEM))
#define ION_HEAP_SYSTEM_CONTIG_MASK ((1 << ION_HEAP_TYPE_SYSTEM_CONTIG))
#define ION_HEAP_CARVEOUT_MASK ((1 << ION_HEAP_TYPE_CARVEOUT))
#define ION_HEAP_TYPE_DMA_MASK ((1 << ION_HEAP_TYPE_DMA))
#define ION_NUM_HEAP_IDS (sizeof(unsigned int) * 8)
/**
* allocation flags - the lower 16 bits are used by core ion, the upper 16
* bits are reserved for use by the heaps themselves.
*/
#define ION_FLAG_CACHED 1 /*
* mappings of this buffer should be
* cached, ion will do cache
* maintenance when the buffer is
* mapped for dma
*/
#define ION_FLAG_CACHED_NEEDS_SYNC 2 /*
* mappings of this buffer will created
* at mmap time, if this is set
* caches must be managed
* manually
*/
/**
* DOC: Ion Userspace API
*
* create a client by opening /dev/ion
* most operations handled via following ioctls
*
*/
/**
* struct ion_allocation_data - metadata passed from userspace for allocations
* @len: size of the allocation
* @align: required alignment of the allocation
* @heap_id_mask: mask of heap ids to allocate from
* @flags: flags passed to heap
* @handle: pointer that will be populated with a cookie to use to
* refer to this allocation
*
* Provided by userspace as an argument to the ioctl
*/
struct ion_allocation_data {
size_t len;
size_t align;
unsigned int heap_id_mask;
unsigned int flags;
ion_user_handle_t handle;
};
/**
* struct ion_fd_data - metadata passed to/from userspace for a handle/fd pair
* @handle: a handle
* @fd: a file descriptor representing that handle
*
* For ION_IOC_SHARE or ION_IOC_MAP userspace populates the handle field with
* the handle returned from ion alloc, and the kernel returns the file
* descriptor to share or map in the fd field. For ION_IOC_IMPORT, userspace
* provides the file descriptor and the kernel returns the handle.
*/
struct ion_fd_data {
ion_user_handle_t handle;
int fd;
};
/**
* struct ion_handle_data - a handle passed to/from the kernel
* @handle: a handle
*/
struct ion_handle_data {
ion_user_handle_t handle;
};
/**
* struct ion_custom_data - metadata passed to/from userspace for a custom ioctl
* @cmd: the custom ioctl function to call
* @arg: additional data to pass to the custom ioctl, typically a user
* pointer to a predefined structure
*
* This works just like the regular cmd and arg fields of an ioctl.
*/
struct ion_custom_data {
unsigned int cmd;
unsigned long arg;
};
#define ION_IOC_MAGIC 'I'
/**
* DOC: ION_IOC_ALLOC - allocate memory
*
* Takes an ion_allocation_data struct and returns it with the handle field
* populated with the opaque handle for the allocation.
*/
#define ION_IOC_ALLOC _IOWR(ION_IOC_MAGIC, 0, \
struct ion_allocation_data)
/**
* DOC: ION_IOC_FREE - free memory
*
* Takes an ion_handle_data struct and frees the handle.
*/
#define ION_IOC_FREE _IOWR(ION_IOC_MAGIC, 1, struct ion_handle_data)
/**
* DOC: ION_IOC_MAP - get a file descriptor to mmap
*
* Takes an ion_fd_data struct with the handle field populated with a valid
* opaque handle. Returns the struct with the fd field set to a file
* descriptor open in the current address space. This file descriptor
* can then be used as an argument to mmap.
*/
#define ION_IOC_MAP _IOWR(ION_IOC_MAGIC, 2, struct ion_fd_data)
/**
* DOC: ION_IOC_SHARE - creates a file descriptor to use to share an allocation
*
* Takes an ion_fd_data struct with the handle field populated with a valid
* opaque handle. Returns the struct with the fd field set to a file
* descriptor open in the current address space. This file descriptor
* can then be passed to another process. The corresponding opaque handle can
* be retrieved via ION_IOC_IMPORT.
*/
#define ION_IOC_SHARE _IOWR(ION_IOC_MAGIC, 4, struct ion_fd_data)
/**
* DOC: ION_IOC_IMPORT - imports a shared file descriptor
*
* Takes an ion_fd_data struct with the fd field populated with a valid file
* descriptor obtained from ION_IOC_SHARE and returns the struct with the handle
* filed set to the corresponding opaque handle.
*/
#define ION_IOC_IMPORT _IOWR(ION_IOC_MAGIC, 5, struct ion_fd_data)
/**
* DOC: ION_IOC_SYNC - syncs a shared file descriptors to memory
*
* Deprecated in favor of using the dma_buf api's correctly (syncing
* will happen automatically when the buffer is mapped to a device).
* If necessary should be used after touching a cached buffer from the cpu,
* this will make the buffer in memory coherent.
*/
#define ION_IOC_SYNC _IOWR(ION_IOC_MAGIC, 7, struct ion_fd_data)
/**
* DOC: ION_IOC_CUSTOM - call architecture specific ion ioctl
*
* Takes the argument of the architecture specific ioctl to call and
* passes appropriate userdata for that ioctl
*/
#define ION_IOC_CUSTOM _IOWR(ION_IOC_MAGIC, 6, struct ion_custom_data)
#endif /* _UAPI_LINUX_ION_H */

211
artifacts/package_sources/tinygrad/extra/dsp/include/msm_ion.h Normal file
View File

@@ -0,0 +1,211 @@
#ifndef _UAPI_MSM_ION_H
#define _UAPI_MSM_ION_H
#include "ion.h"
enum msm_ion_heap_types {
ION_HEAP_TYPE_MSM_START = ION_HEAP_TYPE_CUSTOM + 1,
ION_HEAP_TYPE_SECURE_DMA = ION_HEAP_TYPE_MSM_START,
ION_HEAP_TYPE_SYSTEM_SECURE,
ION_HEAP_TYPE_HYP_CMA,
/*
* if you add a heap type here you should also add it to
* heap_types_info[] in msm_ion.c
*/
};
/**
* These are the only ids that should be used for Ion heap ids.
* The ids listed are the order in which allocation will be attempted
* if specified. Don't swap the order of heap ids unless you know what
* you are doing!
* Id's are spaced by purpose to allow new Id's to be inserted in-between (for
* possible fallbacks)
*/
enum ion_heap_ids {
INVALID_HEAP_ID = -1,
ION_CP_MM_HEAP_ID = 8,
ION_SECURE_HEAP_ID = 9,
ION_SECURE_DISPLAY_HEAP_ID = 10,
ION_CP_MFC_HEAP_ID = 12,
ION_CP_WB_HEAP_ID = 16, /* 8660 only */
ION_CAMERA_HEAP_ID = 20, /* 8660 only */
ION_SYSTEM_CONTIG_HEAP_ID = 21,
ION_ADSP_HEAP_ID = 22,
ION_PIL1_HEAP_ID = 23, /* Currently used for other PIL images */
ION_SF_HEAP_ID = 24,
ION_SYSTEM_HEAP_ID = 25,
ION_PIL2_HEAP_ID = 26, /* Currently used for modem firmware images */
ION_QSECOM_HEAP_ID = 27,
ION_AUDIO_HEAP_ID = 28,
ION_MM_FIRMWARE_HEAP_ID = 29,
ION_HEAP_ID_RESERVED = 31 /** Bit reserved for ION_FLAG_SECURE flag */
};
/*
* The IOMMU heap is deprecated! Here are some aliases for backwards
* compatibility:
*/
#define ION_IOMMU_HEAP_ID ION_SYSTEM_HEAP_ID
#define ION_HEAP_TYPE_IOMMU ION_HEAP_TYPE_SYSTEM
enum ion_fixed_position {
NOT_FIXED,
FIXED_LOW,
FIXED_MIDDLE,
FIXED_HIGH,
};
enum cp_mem_usage {
VIDEO_BITSTREAM = 0x1,
VIDEO_PIXEL = 0x2,
VIDEO_NONPIXEL = 0x3,
DISPLAY_SECURE_CP_USAGE = 0x4,
CAMERA_SECURE_CP_USAGE = 0x5,
MAX_USAGE = 0x6,
UNKNOWN = 0x7FFFFFFF,
};
/**
* Flags to be used when allocating from the secure heap for
* content protection
*/
#define ION_FLAG_CP_TOUCH (1 << 17)
#define ION_FLAG_CP_BITSTREAM (1 << 18)
#define ION_FLAG_CP_PIXEL (1 << 19)
#define ION_FLAG_CP_NON_PIXEL (1 << 20)
#define ION_FLAG_CP_CAMERA (1 << 21)
#define ION_FLAG_CP_HLOS (1 << 22)
#define ION_FLAG_CP_HLOS_FREE (1 << 23)
#define ION_FLAG_CP_SEC_DISPLAY (1 << 25)
#define ION_FLAG_CP_APP (1 << 26)
/**
* Flag to allow non continguous allocation of memory from secure
* heap
*/
#define ION_FLAG_ALLOW_NON_CONTIG (1 << 24)
/**
* Flag to use when allocating to indicate that a heap is secure.
*/
#define ION_FLAG_SECURE (1 << ION_HEAP_ID_RESERVED)
/**
* Flag for clients to force contiguous memort allocation
*
* Use of this flag is carefully monitored!
*/
#define ION_FLAG_FORCE_CONTIGUOUS (1 << 30)
/*
* Used in conjunction with heap which pool memory to force an allocation
* to come from the page allocator directly instead of from the pool allocation
*/
#define ION_FLAG_POOL_FORCE_ALLOC (1 << 16)
#define ION_FLAG_POOL_PREFETCH (1 << 27)
/**
* Deprecated! Please use the corresponding ION_FLAG_*
*/
#define ION_SECURE ION_FLAG_SECURE
#define ION_FORCE_CONTIGUOUS ION_FLAG_FORCE_CONTIGUOUS
/**
* Macro should be used with ion_heap_ids defined above.
*/
#define ION_HEAP(bit) (1 << (bit))
#define ION_ADSP_HEAP_NAME "adsp"
#define ION_SYSTEM_HEAP_NAME "system"
#define ION_VMALLOC_HEAP_NAME ION_SYSTEM_HEAP_NAME
#define ION_KMALLOC_HEAP_NAME "kmalloc"
#define ION_AUDIO_HEAP_NAME "audio"
#define ION_SF_HEAP_NAME "sf"
#define ION_MM_HEAP_NAME "mm"
#define ION_CAMERA_HEAP_NAME "camera_preview"
#define ION_IOMMU_HEAP_NAME "iommu"
#define ION_MFC_HEAP_NAME "mfc"
#define ION_WB_HEAP_NAME "wb"
#define ION_MM_FIRMWARE_HEAP_NAME "mm_fw"
#define ION_PIL1_HEAP_NAME "pil_1"
#define ION_PIL2_HEAP_NAME "pil_2"
#define ION_QSECOM_HEAP_NAME "qsecom"
#define ION_SECURE_HEAP_NAME "secure_heap"
#define ION_SECURE_DISPLAY_HEAP_NAME "secure_display"
#define ION_SET_CACHED(__cache) (__cache | ION_FLAG_CACHED)
#define ION_SET_UNCACHED(__cache) (__cache & ~ION_FLAG_CACHED)
#define ION_IS_CACHED(__flags) ((__flags) & ION_FLAG_CACHED)
/* struct ion_flush_data - data passed to ion for flushing caches
*
* @handle: handle with data to flush
* @fd: fd to flush
* @vaddr: userspace virtual address mapped with mmap
* @offset: offset into the handle to flush
* @length: length of handle to flush
*
* Performs cache operations on the handle. If p is the start address
* of the handle, p + offset through p + offset + length will have
* the cache operations performed
*/
struct ion_flush_data {
ion_user_handle_t handle;
int fd;
void *vaddr;
unsigned int offset;
unsigned int length;
};
struct ion_prefetch_regions {
unsigned int vmid;
size_t *sizes;
unsigned int nr_sizes;
};
struct ion_prefetch_data {
int heap_id;
unsigned long len;
/* Is unsigned long bad? 32bit compiler vs 64 bit compiler*/
struct ion_prefetch_regions *regions;
unsigned int nr_regions;
};
#define ION_IOC_MSM_MAGIC 'M'
/**
* DOC: ION_IOC_CLEAN_CACHES - clean the caches
*
* Clean the caches of the handle specified.
*/
#define ION_IOC_CLEAN_CACHES _IOWR(ION_IOC_MSM_MAGIC, 0, \
struct ion_flush_data)
/**
* DOC: ION_IOC_INV_CACHES - invalidate the caches
*
* Invalidate the caches of the handle specified.
*/
#define ION_IOC_INV_CACHES _IOWR(ION_IOC_MSM_MAGIC, 1, \
struct ion_flush_data)
/**
* DOC: ION_IOC_CLEAN_INV_CACHES - clean and invalidate the caches
*
* Clean and invalidate the caches of the handle specified.
*/
#define ION_IOC_CLEAN_INV_CACHES _IOWR(ION_IOC_MSM_MAGIC, 2, \
struct ion_flush_data)
#define ION_IOC_PREFETCH _IOWR(ION_IOC_MSM_MAGIC, 3, \
struct ion_prefetch_data)
#define ION_IOC_DRAIN _IOWR(ION_IOC_MSM_MAGIC, 4, \
struct ion_prefetch_data)
#endif

1038
artifacts/package_sources/tinygrad/extra/dsp/include/remote_default.h Normal file
View File

File diff suppressed because it is too large Load Diff

326
artifacts/package_sources/tinygrad/extra/dsp/invoke_bug.py Normal file
View File

@@ -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)

279
artifacts/package_sources/tinygrad/extra/dsp/invoke_bug_2.py Normal file
View File

@@ -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)

27
artifacts/package_sources/tinygrad/extra/dsp/opt.py Normal file
View File

@@ -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)

79
artifacts/package_sources/tinygrad/extra/dsp/preload.c Normal file
View File

@@ -0,0 +1,79 @@
__attribute__((constructor))
void preload_init() {
Py_Initialize();
PyRun_SimpleString("print('hello from c'); import extra.dsp.hook");
}
#define _GNU_SOURCE // Must be defined before any includes for RTLD_NEXT
#include <stdio.h>
#include <dlfcn.h>
#include <Python.h> // Include Python header
//#include <sys/ioctl.h>
// Define the original ioctl function pointer
static int (*real_ioctl)(int fd, unsigned long request, void *arg) = NULL;
// Our custom ioctl hook
int ioctl(int fd, unsigned long request, void *arg) {
// Initialize the real ioctl function pointer on first call
if (!real_ioctl) {
real_ioctl = dlsym(RTLD_NEXT, "ioctl");
if (!real_ioctl) {
fprintf(stderr, "Error: Could not find real ioctl\n");
return -1;
}
}
// Log the call
//printf("Hooked ioctl: tid=%d fd=%d, request=0x%lx, arg=%p\n", gettid(), fd, request, arg);
// Call a Python function from extra.dsp.hook
PyObject *pName, *pModule, *pFunc, *pArgs, *pValue;
PyGILState_STATE gstate;
// Ensure the GIL is held (required for Python calls in multi-threaded apps)
//gstate = PyGILState_Ensure();
// Import the module
pName = PyUnicode_FromString("extra.dsp.hook");
pModule = PyImport_Import(pName);
Py_DECREF(pName);
// Call the original ioctl
int ret = real_ioctl(fd, request, arg);
if (pModule != NULL) {
// Get the function (assume it’s called "handle_ioctl")
pFunc = PyObject_GetAttrString(pModule, "handle_ioctl");
if (pFunc && PyCallable_Check(pFunc)) {
// Create arguments tuple (fd, request, arg, ret)
pArgs = PyTuple_Pack(4,
PyLong_FromLong(fd),
PyLong_FromUnsignedLong(request),
PyLong_FromVoidPtr(arg),
PyLong_FromLong(ret));
pValue = PyObject_CallObject(pFunc, pArgs);
Py_DECREF(pArgs);
if (pValue != NULL) {
Py_DECREF(pValue);
} else {
PyErr_Print(); // Print Python error if call fails
}
Py_DECREF(pFunc);
} else {
if (PyErr_Occurred()) PyErr_Print();
fprintf(stderr, "Cannot find function 'handle_ioctl'\n");
}
Py_DECREF(pModule);
} else {
PyErr_Print();
fprintf(stderr, "Failed to load 'extra.dsp.hook'\n");
}
// Release the GIL
//PyGILState_Release(gstate);
return ret;
}

152
artifacts/package_sources/tinygrad/extra/dsp/run.py Executable file
View File

@@ -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)

312
artifacts/package_sources/tinygrad/extra/dsp/run_3.py Executable file
View File

@@ -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)

11
artifacts/package_sources/tinygrad/extra/dsp/snpe.sh Executable file
View File

@@ -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

715
artifacts/package_sources/tinygrad/extra/dsp/snpe_logs/dlc_info_2 Normal file
View File

@@ -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
-----------------------------------------------------------------------------------------------------------------------------------------

131
artifacts/package_sources/tinygrad/extra/dsp/snpe_logs/high_perf_2 Normal file
View File

@@ -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

21
artifacts/package_sources/tinygrad/extra/dsp/snpe_logs/parse.py Normal file
View File

@@ -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")

303
artifacts/package_sources/tinygrad/extra/export_model.py Normal file
View File

@@ -0,0 +1,303 @@
from typing import Tuple, Dict, List, Optional
from tinygrad.dtype import DType, dtypes, AddrSpace
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[2].arg
cargs = [name_of(bu, i == 0) for i, bu in enumerate(arg_uops)] + list(info.vars)
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: NUM_CPU_THREADS=1, since export does not support threading
with Context(JIT=2, NUM_CPU_THREADS=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.PARAM and var.addrspace is AddrSpace.ALU and var.arg.name is not None:
if var not in symbolic_vars:
symbolic_vars[var] = var.expr
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 \
any(s.op is Ops.PARAM and s.addrspace is AddrSpace.ALU for s in dim.src) and any(s.op is Ops.CONST for s in dim.src):
name, val = dim.src if dim.src[1].op is Ops.CONST else reversed(dim.src)
global_size[j] = f"_{name.expr}[0] + {val.val}"
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

16
artifacts/package_sources/tinygrad/extra/f16_decompress.py Normal file
View File

@@ -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).bool()
exponent = bit_extract(x, 14, 10).float()
fraction = bit_extract(x, 9, 0).float()
return sign.where(-1, 1) * exponent.bool().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()

101
artifacts/package_sources/tinygrad/extra/fp8/fp8_linear.py Normal file
View File

@@ -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)
out_idx = UOp.range(OUT, 3)
batch_idx = UOp.range(output.size//SEQ//OUT, 1)
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)).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.unshard(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

2
artifacts/package_sources/tinygrad/extra/gemm/.gitignore vendored Normal file
View File

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

501
artifacts/package_sources/tinygrad/extra/gemm/amd_asm_matmul.py Normal file
View File

@@ -0,0 +1,501 @@
# 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_size = max(LDS_SIZE, 65536//getenv("LIMIT_OCC", 65536))
lds = UOp.placeholder((lds_size,), dtypes.uint8, 0, AddrSpace.LOCAL)
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.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()

118
artifacts/package_sources/tinygrad/extra/gemm/amd_copy_matmul.py Normal file
View File

@@ -0,0 +1,118 @@
from tinygrad import Device, UOp, getenv
from tinygrad.uop.ops import AxisType, KernelInfo
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, slot=0, addrspace=AddrSpace.LOCAL)
B_local = UOp.placeholder((BLOCK_N, BLOCK_K) if use_wmma else (BLOCK_K, BLOCK_N), b.dtype, 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])
# NOTE: no explicit barrier needed, the AFTER on the LOCAL buffers implies it in late codegen
A_local, B_local = A_local.after(A_store, B_store), B_local.after(A_store, B_store)
# -- 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(buffer=False)))
if use_wmma:
k = UOp.range(BLOCK_K // WMMA_K, 101, AxisType.REDUCE)
tile_m = UOp.range(TM // WMMA_ACC, 200)
tile_n = UOp.range(TN, 201)
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.wmma(a_frag, b_frag, acc_frag.after(k), (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 (the barrier at the end of the loop is implied by the LOCAL buffers stored and loaded in the loop)
acc = acc.after(acc_store.end(k).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)

50
artifacts/package_sources/tinygrad/extra/gemm/amd_matmul.py Normal file
View File

@@ -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.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()}")

143
artifacts/package_sources/tinygrad/extra/gemm/amd_seb/kernel3_registers.cpp Normal file
View File

@@ -0,0 +1,143 @@
typedef long unsigned int size_t;
extern "C" __attribute__((device, const)) size_t __ockl_get_local_id(unsigned int);
extern "C" __attribute__((device, const)) size_t __ockl_get_group_id(unsigned int);
struct Dim3 { size_t x, y, z; };
#define __shared__ __attribute__((shared, aligned(16)))
__attribute__((device)) inline void __syncthreads() {
__builtin_amdgcn_fence(__ATOMIC_RELEASE, "workgroup");
__builtin_amdgcn_s_barrier();
__builtin_amdgcn_fence(__ATOMIC_ACQUIRE, "workgroup");
}
#define BLOCK_SIZE 256
extern "C" __attribute__((global)) void __attribute__((amdgpu_flat_work_group_size(1, BLOCK_SIZE)))
kernel3_registers(float *a, float *b, float *c)
{
constexpr int N = 4096;
constexpr float alpha = 1.0;
constexpr float beta = 0.0;
const Dim3 blockIdx{ __ockl_get_group_id(0), __ockl_get_group_id(1), __ockl_get_group_id(2) };
const Dim3 threadIdx{ __ockl_get_local_id(0), __ockl_get_local_id(1), __ockl_get_local_id(2) };
// Block Tile size
constexpr int BN = 128;
constexpr int BM = 128;
// Number of Row or column we read per batch
constexpr int BK = 8;
// Thread Tile size
constexpr int TN = 4;
constexpr int TM = 4;
constexpr int nbWaves = BLOCK_SIZE / 32;
// Wave Tile size
constexpr int WN = 64;
constexpr int WM = BN * BM / nbWaves / WN;
// Number of wave on X & Y axis in the Block tile
constexpr int nbWaveX = BN / WN;
constexpr int nbWaveY = BM / WM;
const int waveIndex = threadIdx.x / 32;
const int waveIdx = waveIndex % nbWaveX;
const int waveIdy = waveIndex / nbWaveX;
const int indexInWave = threadIdx.x % 32;
// A wave is a block of 8x4 of the output matrix
constexpr int nbThreadXPerWave = 8;
constexpr int nbThreadYPerWave = 4;
// Thread coordinates in Wave
const int idxInWave = indexInWave % nbThreadXPerWave;
const int idyInWave = indexInWave / nbThreadXPerWave;
constexpr int nbIterWaveN = WN / (nbThreadXPerWave * TN);
constexpr int nbIterWaveM = WM / (nbThreadYPerWave * TM);
// Wave Sub-tile size
constexpr int SUBWN = WN / nbIterWaveN;
constexpr int SUBWM = WM / nbIterWaveM;
// Thread mapping to read BKxBN block from A
int rAIdx = threadIdx.x % BK;
int rAIdy = threadIdx.x / BK;
// Thread mapping to read BNxBK block from B
int rBIdx = threadIdx.x % BN;
int rBIdy = threadIdx.x / BN;
constexpr int strideReadB = BLOCK_SIZE / BN;
constexpr int strideReadA = BLOCK_SIZE / BK;
constexpr int nbReadsB = BN * BK / BLOCK_SIZE;
constexpr int nbReadsA = BM * BK / BLOCK_SIZE;
float A_col[nbIterWaveM * TM];
float B_row[nbIterWaveN * TN];
__shared__ float As[BK][BM];
__shared__ float Bs[BK][BN];
float c_regs[TM * nbIterWaveM * TN * nbIterWaveN] = {0.0f};
// Iteration over BK blocks.
for (int kId = 0; kId < N; kId += BK) {
__syncthreads();
// We populate the Shared Memory with Ks row and columns
for (int i = 0; i < nbReadsB; i++) {
int index_x = BN * blockIdx.x + rBIdx;
int index_y = rBIdy + i * strideReadB + kId;
Bs[index_y % BK][index_x % BN] = b[N * index_y + index_x];
}
for (int i = 0; i < nbReadsA; i++) {
int index_x = rAIdx + kId;
int index_y = BM * blockIdx.y + rAIdy + i * strideReadA;
As[(index_x % BK)][(index_y % BM)] = a[N * index_y + index_x];
}
__syncthreads();
for (int k = 0; k < BK; k++) {
// we cache A & B for the entire Wave tile
for (int iterWave = 0; iterWave < nbIterWaveN; iterWave++) {
for (int i = 0; i < TN; i++) {
int index = waveIdx * WN + iterWave * SUBWN + TN * idxInWave + i;
B_row[iterWave * TN + i] = Bs[k][index];
}
}
for (int iterWave = 0; iterWave < nbIterWaveM; iterWave++) {
for (int i = 0; i < TM; i++) {
int index = waveIdy * WM + iterWave * SUBWM + TM * idyInWave + i;
A_col[iterWave * TM + i] = As[k][index];
}
}
// we accumulate to C_regs
for (int iterWaveM = 0; iterWaveM < nbIterWaveM; iterWaveM++) {
for (int iterWaveN = 0; iterWaveN < nbIterWaveN; iterWaveN++) {
for (int yt = 0; yt < TM; yt++) {
for (int xt = 0; xt < TN; xt++) {
const int x = iterWaveN * TN + xt;
const int y = iterWaveM * TM + yt;
c_regs[y * TN * nbIterWaveN + x] += A_col[y] * B_row[x];
}
}
}
}
}
}
for (int iterWaveM = 0; iterWaveM < nbIterWaveM; iterWaveM++) {
for (int iterWaveN = 0; iterWaveN < nbIterWaveN; iterWaveN++) {
int xOut = blockIdx.x * BN + waveIdx * WN + iterWaveN * SUBWN + TN * idxInWave;
int yOut = blockIdx.y * BM + waveIdy * WM + iterWaveM * SUBWM + TM * idyInWave;
for (int yt = 0; yt < TM; yt++) {
for (int xt = 0; xt < TN; xt++) {
int indexC = N * (yOut + yt) + xOut + xt;
c[indexC] = beta * c[indexC] + alpha * c_regs[TN * nbIterWaveN * (iterWaveM * TM + yt) + (iterWaveN * TN + xt)];
}
}
}
}
}

172
artifacts/package_sources/tinygrad/extra/gemm/amd_seb/kernel4_gmem_df.cpp Normal file
View File

@@ -0,0 +1,172 @@
typedef long unsigned int size_t;
extern "C" __attribute__((device, const)) size_t __ockl_get_local_id(unsigned int);
extern "C" __attribute__((device, const)) size_t __ockl_get_group_id(unsigned int);
struct Dim3 { size_t x, y, z; };
#define __shared__ __attribute__((shared, aligned(16)))
__attribute__((device)) inline void __syncthreads() {
__builtin_amdgcn_fence(__ATOMIC_RELEASE, "workgroup");
__builtin_amdgcn_s_barrier();
__builtin_amdgcn_fence(__ATOMIC_ACQUIRE, "workgroup");
}
#define BLOCK_SIZE 256
extern "C" __attribute__((global)) void __attribute__((amdgpu_flat_work_group_size(1, BLOCK_SIZE)))
kernel4_gmem_db(float *a, float *b, float *c)
{
constexpr int N = 4096;
constexpr float alpha = 1.0;
constexpr float beta = 0.0;
const Dim3 blockIdx{ __ockl_get_group_id(0), __ockl_get_group_id(1), __ockl_get_group_id(2) };
const Dim3 threadIdx{ __ockl_get_local_id(0), __ockl_get_local_id(1), __ockl_get_local_id(2) };
// Block Tile size
constexpr int BN = 128;
constexpr int BM = 128;
// Number of Row or column we read per batch
constexpr int BK = 8;
// Thread Tile size
constexpr int TN = 4;
constexpr int TM = 4;
constexpr int nbWaves = BLOCK_SIZE / 32;
// Wave Tile size
constexpr int WN = 64;
constexpr int WM = BN * BM / nbWaves / WN;
// Number of wave on X & Y axis in the Block tile
constexpr int nbWaveX = BN / WN;
constexpr int nbWaveY = BM / WM;
const int waveIndex = threadIdx.x / 32;
const int waveIdx = waveIndex % nbWaveX;
const int waveIdy = waveIndex / nbWaveX;
const int indexInWave = threadIdx.x % 32;
// A wave is a block of 8x4 of the output matrix
constexpr int nbThreadXPerWave = 8;
constexpr int nbThreadYPerWave = 4;
// Thread coordinates in Wave
const int idxInWave = indexInWave % nbThreadXPerWave;
const int idyInWave = indexInWave / nbThreadXPerWave;
constexpr int nbIterWaveN = WN / (nbThreadXPerWave * TN);
constexpr int nbIterWaveM = WM / (nbThreadYPerWave * TM);
// Wave Sub-tile size
constexpr int SUBWN = WN / nbIterWaveN;
constexpr int SUBWM = WM / nbIterWaveM;
// Thread mapping to read BKxBN block from A
int rAIdx = threadIdx.x % BK;
int rAIdy = threadIdx.x / BK;
// Thread mapping to read BNxBK block from B
int rBIdx = threadIdx.x % BN;
int rBIdy = threadIdx.x / BN;
constexpr int strideReadB = BLOCK_SIZE / BN;
constexpr int strideReadA = BLOCK_SIZE / BK;
constexpr int nbReadsB = BN * BK / BLOCK_SIZE;
constexpr int nbReadsA = BM * BK / BLOCK_SIZE;
float A_col[nbIterWaveM * TM];
float B_row[nbIterWaveN * TN];
__shared__ float As[BK][BM];
__shared__ float Bs[BK][BN];
float c_regs[TM * nbIterWaveM * TN * nbIterWaveN] = {0.0f};
for (int i = 0; i < nbReadsB; i++) {
int index_x = BN * blockIdx.x + rBIdx;
int index_y = rBIdy + i * strideReadB;
Bs[index_y % BK][index_x % BN] = b[N * index_y + index_x];
}
for (int i = 0; i < nbReadsA; i++) {
int index_x = rAIdx;
int index_y = BM * blockIdx.y + rAIdy + i * strideReadA;
As[(index_x % BK)][(index_y % BM)] = a[N * index_y + index_x];
}
__syncthreads();
// Iteration over BK blocks.
for (int kId = 0; kId < N; kId += BK) {
float regA[nbReadsA];
float regB[nbReadsB];
if (kId < N - BK) {
// We populate the Shared Memory with Ks row and columns
for (int i = 0; i < nbReadsB; i++) {
int index_x = BN * blockIdx.x + rBIdx;
int index_y = rBIdy + i * strideReadB + kId + BK;
regB[i] = b[N * index_y + index_x];
}
for (int i = 0; i < nbReadsA; i++) {
int index_x = rAIdx + kId + BK;
int index_y = BM * blockIdx.y + rAIdy + i * strideReadA;
regA[i] = a[N * index_y + index_x];
}
}
for (int k = 0; k < BK; k++) {
// we cache A & B for the entire Wave tile
for (int iterWave = 0; iterWave < nbIterWaveN; iterWave++) {
for (int i = 0; i < TN; i++) {
int index = waveIdx * WN + iterWave * SUBWN + TN * idxInWave + i;
B_row[iterWave * TN + i] = Bs[k][index];
}
}
for (int iterWave = 0; iterWave < nbIterWaveM; iterWave++) {
for (int i = 0; i < TM; i++) {
int index = waveIdy * WM + iterWave * SUBWM + TM * idyInWave + i;
A_col[iterWave * TM + i] = As[k][index];
}
}
// we accumulate to C_regs
for (int iterWaveM = 0; iterWaveM < nbIterWaveM; iterWaveM++) {
for (int iterWaveN = 0; iterWaveN < nbIterWaveN; iterWaveN++) {
for (int yt = 0; yt < TM; yt++) {
for (int xt = 0; xt < TN; xt++) {
const int x = iterWaveN * TN + xt;
const int y = iterWaveM * TM + yt;
c_regs[y * TN * nbIterWaveN + x] += A_col[y] * B_row[x];
}
}
}
}
}
__syncthreads();
if (kId < N - BK) {
for (int i = 0; i < nbReadsB; i++) {
int index_x = BN * blockIdx.x + rBIdx;
int index_y = rBIdy + i * strideReadB + kId + BK;
Bs[index_y % BK][index_x % BN] = regB[i]; // row
}
for (int i = 0; i < nbReadsA; i++) {
int index_x = rAIdx + kId + BK;
int index_y = BM * blockIdx.y + rAIdy + i * strideReadA;
As[(index_x % BK)][(index_y % BM)] = regA[i];
}
__syncthreads();
}
}
for (int iterWaveM = 0; iterWaveM < nbIterWaveM; iterWaveM++) {
for (int iterWaveN = 0; iterWaveN < nbIterWaveN; iterWaveN++) {
int xOut = blockIdx.x * BN + waveIdx * WN + iterWaveN * SUBWN + TN * idxInWave;
int yOut = blockIdx.y * BM + waveIdy * WM + iterWaveM * SUBWM + TM * idyInWave;
for (int yt = 0; yt < TM; yt++) {
for (int xt = 0; xt < TN; xt++) {
int indexC = N * (yOut + yt) + xOut + xt;
c[indexC] = beta * c[indexC] + alpha * c_regs[TN * nbIterWaveN * (iterWaveM * TM + yt) + (iterWaveN * TN + xt)];
}
}
}
}
}

172
artifacts/package_sources/tinygrad/extra/gemm/amd_seb/kernel5_lds_optim.cpp Normal file
View File

@@ -0,0 +1,172 @@
typedef long unsigned int size_t;
extern "C" __attribute__((device, const)) size_t __ockl_get_local_id(unsigned int);
extern "C" __attribute__((device, const)) size_t __ockl_get_group_id(unsigned int);
struct Dim3 { size_t x, y, z; };
#define __shared__ __attribute__((shared, aligned(16)))
__attribute__((device)) inline void __syncthreads() {
__builtin_amdgcn_fence(__ATOMIC_RELEASE, "workgroup");
__builtin_amdgcn_s_barrier();
__builtin_amdgcn_fence(__ATOMIC_ACQUIRE, "workgroup");
}
#define BLOCK_SIZE 128
extern "C" __attribute__((global)) void __attribute__((amdgpu_flat_work_group_size(1, BLOCK_SIZE)))
kernel5_lds_optim(float *a, float *b, float *c)
{
constexpr int N = 4096;
constexpr float alpha = 1.0;
constexpr float beta = 0.0;
const Dim3 blockIdx{ __ockl_get_group_id(0), __ockl_get_group_id(1), __ockl_get_group_id(2) };
const Dim3 threadIdx{ __ockl_get_local_id(0), __ockl_get_local_id(1), __ockl_get_local_id(2) };
// Block Tile size
constexpr int BN = 128;
constexpr int BM = 128;
// Number of Row or column we read per batch
constexpr int BK = 8;
// Thread Tile size
constexpr int TN = 4;
constexpr int TM = 4;
constexpr int nbWaves = BLOCK_SIZE / 32;
// Wave Tile size
constexpr int WN = 128;
constexpr int WM = BN * BM / nbWaves / WN;
// Number of wave on X & Y axis in the Block tile
constexpr int nbWaveX = BN / WN;
constexpr int nbWaveY = BM / WM;
const int waveIndex = threadIdx.x / 32;
const int waveIdx = waveIndex % nbWaveX;
const int waveIdy = waveIndex / nbWaveX;
const int indexInWave = threadIdx.x % 32;
// A wave is a block of 8x4 of the output matrix
constexpr int nbThreadXPerWave = 8;
constexpr int nbThreadYPerWave = 4;
// Thread coordinates in Wave
const int idxInWave = indexInWave % nbThreadXPerWave;
const int idyInWave = indexInWave / nbThreadXPerWave;
constexpr int nbIterWaveN = WN / (nbThreadXPerWave * TN);
constexpr int nbIterWaveM = WM / (nbThreadYPerWave * TM);
// Wave Sub-tile size
constexpr int SUBWN = WN / nbIterWaveN;
constexpr int SUBWM = WM / nbIterWaveM;
// Thread mapping to read BKxBN block from A
int rAIdx = threadIdx.x % BK;
int rAIdy = threadIdx.x / BK;
// Thread mapping to read BNxBK block from B
int rBIdx = threadIdx.x % BN;
int rBIdy = threadIdx.x / BN;
constexpr int strideReadB = BLOCK_SIZE / BN;
constexpr int strideReadA = BLOCK_SIZE / BK;
constexpr int nbReadsB = BN * BK / BLOCK_SIZE;
constexpr int nbReadsA = BM * BK / BLOCK_SIZE;
float A_col[nbIterWaveM * TM];
float B_row[nbIterWaveN * TN];
__shared__ float As[BK][BM+4]; // 4 padding to avoid bank conflicts
__shared__ float Bs[BK][BN];
float c_regs[TM * nbIterWaveM * TN * nbIterWaveN] = {0.0f};
// initial copy into shared memory
for (int i = 0; i < nbReadsB; i++) {
int index_x = BN * blockIdx.x + rBIdx;
int index_y = rBIdy + i * strideReadB;
Bs[index_y % BK][index_x % BN] = b[N * index_y + index_x];
}
for (int i = 0; i < nbReadsA; i++) {
int index_x = rAIdx;
int index_y = BM * blockIdx.y + rAIdy + i * strideReadA;
As[(index_x % BK)][(index_y % BM)] = a[N * index_y + index_x];
}
__syncthreads();
// Iteration over BK blocks.
for (int kId = 0; kId < N; kId += BK) {
float regA[nbReadsA];
float regB[nbReadsB];
if (kId < N - BK) {
// We populate the Shared Memory with Ks row and columns
for (int i = 0; i < nbReadsB; i++) {
int index_x = BN * blockIdx.x + rBIdx;
int index_y = rBIdy + i * strideReadB + kId + BK;
regB[i] = b[N * index_y + index_x];
}
for (int i = 0; i < nbReadsA; i++) {
int index_x = rAIdx + kId + BK;
int index_y = BM * blockIdx.y + rAIdy + i * strideReadA;
regA[i] = a[N * index_y + index_x];
}
}
for (int k = 0; k < BK; k++) {
// we cache A & B for the entire Wave tile
for (int iterWave = 0; iterWave < nbIterWaveN; iterWave++) {
for (int i = 0; i < TN; i++) {
int index = waveIdx * WN + iterWave * SUBWN + TN * idxInWave + i;
B_row[iterWave * TN + i] = Bs[k][index];
}
}
for (int iterWave = 0; iterWave < nbIterWaveM; iterWave++) {
for (int i = 0; i < TM; i++) {
int index = waveIdy * WM + iterWave * SUBWM + TM * idyInWave + i;
A_col[iterWave * TM + i] = As[k][index];
}
}
// we accumulate to C_regs
for (int iterWaveM = 0; iterWaveM < nbIterWaveM; iterWaveM++) {
for (int iterWaveN = 0; iterWaveN < nbIterWaveN; iterWaveN++) {
for (int yt = 0; yt < TM; yt++) {
for (int xt = 0; xt < TN; xt++) {
const int x = iterWaveN * TN + xt;
const int y = iterWaveM * TM + yt;
c_regs[y * TN * nbIterWaveN + x] += A_col[y] * B_row[x];
}
}
}
}
}
__syncthreads();
if (kId < N - BK) {
for (int i = 0; i < nbReadsB; i++) {
int index_x = BN * blockIdx.x + rBIdx;
int index_y = rBIdy + i * strideReadB + kId + BK;
Bs[index_y % BK][index_x % BN] = regB[i]; // row
}
for (int i = 0; i < nbReadsA; i++) {
int index_x = rAIdx + kId + BK;
int index_y = BM * blockIdx.y + rAIdy + i * strideReadA;
As[(index_x % BK)][(index_y % BM)] = regA[i];
}
__syncthreads();
}
}
for (int iterWaveM = 0; iterWaveM < nbIterWaveM; iterWaveM++) {
for (int iterWaveN = 0; iterWaveN < nbIterWaveN; iterWaveN++) {
int xOut = blockIdx.x * BN + waveIdx * WN + iterWaveN * SUBWN + TN * idxInWave;
int yOut = blockIdx.y * BM + waveIdy * WM + iterWaveM * SUBWM + TM * idyInWave;
for (int yt = 0; yt < TM; yt++) {
for (int xt = 0; xt < TN; xt++) {
int indexC = N * (yOut + yt) + xOut + xt;
c[indexC] = beta * c[indexC] + alpha * c_regs[TN * nbIterWaveN * (iterWaveM * TM + yt) + (iterWaveN * TN + xt)];
}
}
}
}
}

2423
artifacts/package_sources/tinygrad/extra/gemm/amd_seb/kernel8_batched_gmem.s Normal file
View File

File diff suppressed because it is too large Load Diff

142
artifacts/package_sources/tinygrad/extra/gemm/amd_uop_matmul.py Normal file
View File

@@ -0,0 +1,142 @@
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.WEAK): 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.WEAK)
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)
# NOTE: no explicit barrier needed, the AFTER on the LOCAL buffers implies it in late codegen
A_local, B_local = A_local.after(A_local_store, B_local_store), B_local.after(A_local_store, B_local_store)
# 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).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)

478
artifacts/package_sources/tinygrad/extra/gemm/cdna_asm_gemm.py Normal file
View File

@@ -0,0 +1,478 @@
import atexit, functools, math, pathlib
from tinygrad import Tensor, Device, dtypes
from tinygrad.dtype import AddrSpace
from tinygrad.uop.ops import UOp, Ops, KernelInfo, AxisType
from tinygrad.renderer import Estimates
from tinygrad.helpers import getenv, all_same, DEBUG, ceildiv
from tinygrad.runtime.support.compiler_amd import HIPCCCompiler
from examples.mlperf.models.flat_llama import FP8_DTYPE, quantize_fp8
from extra.llama_kernels.quantize_mxfp4 import quantize_mxfp4
TILE_M, TILE_N, TILE_K = 256, 256, 64
# ** FP8 GEMM custom kernel
@functools.cache
def custom_hk_fp8_gemm(C:UOp, A:UOp, B:UOp, *args:UOp, dname:str, scale_mode:int=3) -> UOp:
# scale_mode: 0=no scale, 1=x only, 2=w only, 3=both
n_scales = (1 if scale_mode & 1 else 0) + (1 if scale_mode & 2 else 0) + (1 if scale_mode & 4 else 0)
scales, extra = args[:n_scales], args[n_scales:]
M, K = A.shape[0]*A.shape[1], A.shape[2]
N, K2 = B.shape[(1 if B.ndim == 3 else 0):]
assert K == K2, f"{A.shape} {B.shape}"
block_size = 256
threads = UOp.special(64 * 8, "lidx0")
workgroups = UOp.special((M // block_size) * (N // block_size), "gidx0")
sink_inputs = (C.base, A.base, B.base) + tuple(s.base for s in scales) + (threads, workgroups)
sink = UOp.sink(*sink_inputs,
arg=KernelInfo(f"hk_fp8_gemm_{M}_{N}_{K}", estimates=Estimates(ops=2*M*N*K, mem=(M*K+N*K)*A.dtype.itemsize+M*N*C.dtype.itemsize)))
kittens_path = pathlib.Path(__file__).parent.parent/"thunder"/"amd"
src = (kittens_path/"gemm_fp8.cpp").read_text()
lib = HIPCCCompiler("gfx950", [f"-I{(kittens_path/'include').as_posix()}", "-std=c++20", "-DKITTENS_CDNA4", "-ffast-math",
"-DHIP_ENABLE_WARP_SYNC_BUILTINS", f"-DGEMM_M={M}", f"-DGEMM_N={N}", f"-DGEMM_K={K}",
f"-DSCALE_MODE={scale_mode}"]).compile_cached(src)
return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=src),
UOp(Ops.BINARY, arg=lib)))
# ** FP8 AtB GEMM custom kernel
@functools.cache
def custom_hk_fp8_atb_gemm(C:UOp, A:UOp, B:UOp, *args:UOp, dname:str, scale_mode:int=5) -> UOp:
# C = A.T @ B, A and B are physically [K, M] and [K, N].
n_scales = (1 if scale_mode & 1 else 0) + (1 if scale_mode & 2 else 0) + (1 if scale_mode & 4 else 0)
scales = args[:n_scales]
K, M = A.shape[0]*A.shape[1], A.shape[2]
K2, N = B.shape[0]*B.shape[1], B.shape[2]
assert K == K2, f"{A.shape} {B.shape}"
block_m, block_n, block_k, num_warps = 256, 256, 128, 8
assert M % block_m == 0 and N % block_n == 0 and K % block_k == 0, f"invalid fp8 atb tile {(block_m, block_n, block_k)} for {(M, N, K)}"
threads = UOp.special(64 * num_warps, "lidx0")
workgroups = UOp.special((M // block_m) * (N // block_n), "gidx0")
sink_inputs = (C.base, A.base, B.base) + tuple(s.base for s in scales) + (threads, workgroups)
sink = UOp.sink(*sink_inputs,
arg=KernelInfo(f"hk_fp8_atb_gemm_{M}_{N}_{K}", estimates=Estimates(ops=2*M*N*K, mem=(M*K+N*K)*A.dtype.itemsize+M*N*C.dtype.itemsize)))
kittens_path = pathlib.Path(__file__).parent.parent/"thunder"/"amd"
src = (kittens_path/"gemm_fp8_atb.cpp").read_text()
lib = HIPCCCompiler("gfx950", [f"-I{(kittens_path/'include').as_posix()}", "-std=c++20", "-DKITTENS_CDNA4", "-ffast-math",
"-DHIP_ENABLE_WARP_SYNC_BUILTINS", f"-DGEMM_M={M}", f"-DGEMM_N={N}", f"-DGEMM_K={K}",
f"-DSCALE_MODE={scale_mode}"]).compile_cached(src)
return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=src),
UOp(Ops.BINARY, arg=lib)))
def hk_fp8_atb_gemm(a:Tensor, b:Tensor, x_scale:Tensor|None=None, g_amax:Tensor|None=None) -> Tensor:
assert a.dtype == b.dtype == FP8_DTYPE, f"expected fp8, got {a.dtype} {b.dtype}"
assert a.ndim == b.ndim == 3 and a.shape[:2] == b.shape[:2], f"{a.shape} {b.shape}"
batch, rows, M = a.shape
N = b.shape[2]
assert M % TILE_M == 0 and N % TILE_N == 0 and (batch * rows) % 128 == 0, \
f"fp8 atb shape {a.shape} {b.shape} must produce (M,N,K) multiples of ({TILE_M},{TILE_N},128)"
is_multi = isinstance(a.device, tuple)
reduce_out = False
if is_multi:
ndev = len(a.device)
if a.uop.axis in (0, 1) or b.uop.axis in (0, 1): inv, out_axis, reduce_out = Tensor.invalids(1, M, N, dtype=dtypes.bfloat16, device=a.device), 0, True
elif b.uop.axis == 2: inv, out_axis = Tensor.invalids(1, M, N // ndev, dtype=dtypes.bfloat16, device=a.device), 2
elif a.uop.axis == 2: inv, out_axis = Tensor.invalids(1, M // ndev, N, dtype=dtypes.bfloat16, device=a.device), 1
else: inv, out_axis, reduce_out = Tensor.invalids(1, M, N, dtype=dtypes.bfloat16, device=a.device), 0, True
out = Tensor(inv.uop.unshard(out_axis), device=a.device)
dname = a.device[0]
else:
out = Tensor.invalids(1, M, N, dtype=dtypes.bfloat16, device=a.device)
dname = a.device
dname = dname.split(":")[0]
scales = tuple(s for s in (x_scale, g_amax) if s is not None)
scale_mode = (1 if x_scale is not None else 0) | (4 if g_amax is not None else 0)
out = Tensor.custom_kernel(out, a, b, *scales, fxn=functools.partial(custom_hk_fp8_atb_gemm, dname=dname, scale_mode=scale_mode))[0]
if reduce_out: out = out.sum(0)
return out.squeeze(0) if out.ndim == 3 else out
# ** MXFP8 GEMM custom kernel
@functools.cache
def custom_hk_mxfp8_gemm(C:UOp, A:UOp, B:UOp, scale_A:UOp, scale_B:UOp, *extra:UOp, dname:str) -> UOp:
# mxfp8 block-scaled gemm: A(M,K) @ B(N,K).T, e8m0 1x32 microscales packed (k_iters,dim) uint32
M, K = A.shape[0]*A.shape[1], A.shape[2]
N, K2 = B.shape[(1 if B.ndim == 3 else 0):]
assert K == K2, f"{A.shape} {B.shape}"
block_size = 256
threads = UOp.special(64 * 8, "lidx0")
workgroups = UOp.special((M // block_size) * (N // block_size), "gidx0")
e_a = extra[0].base if len(extra) >= 1 else scale_A.base
e_b = extra[1].base if len(extra) >= 2 else scale_B.base
sink_inputs = (C.base, A.base, B.base, scale_A.base, scale_B.base, e_a, e_b, threads, workgroups)
sink = UOp.sink(*sink_inputs,
arg=KernelInfo(f"hk_mxfp8_gemm_{M}_{N}_{K}", estimates=Estimates(ops=2*M*N*K, mem=(M*K+N*K)*A.dtype.itemsize+M*N*C.dtype.itemsize)))
kittens_path = pathlib.Path(__file__).parent.parent/"thunder"/"amd"
src = (kittens_path/"gemm_mxfp8.cpp").read_text()
lib = HIPCCCompiler("gfx950", [f"-I{(kittens_path/'include').as_posix()}", "-std=c++20", "-DKITTENS_CDNA4", "-ffast-math",
"-DHIP_ENABLE_WARP_SYNC_BUILTINS", f"-DGEMM_M={M}", f"-DGEMM_N={N}", f"-DGEMM_K={K}"]).compile_cached(src)
return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=src),
UOp(Ops.BINARY, arg=lib)))
# ** MXFP4 GEMM custom kernel
@functools.cache
def custom_mxfp4_gemm(C:UOp, A:UOp, B:UOp, scale_a:UOp, scale_b:UOp, *extra:UOp, tile_m:int, tile_n:int) -> UOp:
from extra.gemm.gemm_mxfp4 import build_kernel
M, half_k = math.prod(A.shape[:-1]), A.shape[-1]
N, half_k_b = math.prod(B.shape[:-1]), B.shape[-1]
K = half_k * 2
assert half_k == half_k_b and math.prod(C.shape[:-1]) == M and C.shape[-1] == N
threads = UOp.special(256, "lidx0")
groups_x, groups_y = UOp.special(ceildiv(N, tile_n), "gidx0"), UOp.special(ceildiv(M, tile_m), "gidx1")
lds = UOp.placeholder((163840,), dtypes.uint8, 0, AddrSpace.LOCAL)
sink = UOp.sink(C.base, A.base, B.base, scale_a.base, scale_b.base, *(x.base for x in extra), lds, threads, groups_x, groups_y,
arg=KernelInfo(f"custom_mxfp4_gemm_{M}_{N}_{K}", estimates=Estimates(ops=2*M*N*K)))
insts = build_kernel(M, N, K, tile_m, tile_n)
return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple(UOp(Ops.INS, arg=x) for x in insts))))
def _mxfp4_gemm_quantized(a_q:Tensor, b_q:Tensor, scale_a:Tensor, scale_b:Tensor) -> Tensor:
M, half_k = a_q.shape
N, half_k_b = b_q.shape
assert half_k == half_k_b
is_multi = isinstance(a_q.device, tuple)
reduce_out = is_multi and (a_q.uop.axis == 1 or b_q.uop.axis == 1)
if not is_multi: out = Tensor.invalids(1, M, N, dtype=dtypes.bfloat16, device=a_q.device)
elif reduce_out: out = Tensor(Tensor.invalids(1, M, N, dtype=dtypes.bfloat16, device=a_q.device).uop.unshard(0), device=a_q.device)
elif a_q.uop.axis == 0:
out = Tensor(Tensor.invalids(1, M//len(a_q.device), N, dtype=dtypes.bfloat16, device=a_q.device).uop.unshard(1), device=a_q.device)
elif b_q.uop.axis == 0:
out = Tensor(Tensor.invalids(1, M, N//len(a_q.device), dtype=dtypes.bfloat16, device=a_q.device).uop.unshard(2), device=a_q.device)
else: out = Tensor.invalids(1, M, N, dtype=dtypes.bfloat16, device=a_q.device)
tile_m, tile_n = next((tm, tn) for tm, tn in ((256, 256), (192, 256), (128, 512)) if M % tm == N % tn == 0)
out = Tensor.custom_kernel(out, a_q, b_q, scale_a, scale_b,
fxn=functools.partial(custom_mxfp4_gemm, tile_m=tile_m, tile_n=tile_n))[0]
if reduce_out: out = out.sum(0)
return out.squeeze(0)
def quantize_mxfp8(x:Tensor) -> tuple[Tensor, Tensor, Tensor]:
# 1x32 block scaling along the last axis
*batch, K = x.shape
scale_K = K // 32
amax = x.detach().float().reshape(*batch, scale_K, 32).abs().max(axis=-1)
e8 = (amax.maximum(1e-38).log2().floor() + 127).clamp(0, 254).cast(dtypes.uint8)
qscale = (127.0 - e8.cast(dtypes.float32)).exp2().reshape(*batch, scale_K, 1).expand(*batch, scale_K, 32).reshape(*batch, K)
x_scaled = x.float() * qscale
x_clamped = x_scaled + (x_scaled.detach().clamp(-448.0, 448.0) - x_scaled.detach()) # STE
packed = mx_pack(e8) if len(batch) == 1 and scale_K % 4 == 0 else None
return x_clamped.cast(FP8_DTYPE), e8, packed
def mx_pack(e8:Tensor) -> Tensor:
rows, scale_K = e8.shape
return e8.reshape(rows, scale_K // 4, 4).bitcast(dtypes.uint32).reshape(rows, scale_K // 4).permute(1, 0).contiguous()
def _mx_block_scale(e8:Tensor) -> Tensor:
# dequant scale 2^(e8-127) broadcast back to element shape
rows, scale_K = e8.shape
return (e8.cast(dtypes.float32) - 127.0).exp2().reshape(rows, scale_K, 1).expand(rows, scale_K, 32).reshape(rows, scale_K*32)
def _mx_block_scale_3d(e8:Tensor) -> Tensor:
# batched (E, rows, scale_K) dequant scale 2^(e8-127) broadcast to (E, rows, scale_K*32)
E, rows, scale_K = e8.shape
return (e8.cast(dtypes.float32) - 127.0).exp2().reshape(E, rows, scale_K, 1).expand(E, rows, scale_K, 32).reshape(E, rows, scale_K*32)
counters = {"used":0, "todos":[]}
def todo(msg:str) -> bool: counters["todos"].append(msg); return False
def _asm_gemm_report():
print(f'asm_gemm: {counters["used"]} used, {len(counters["todos"])} not used')
if DEBUG >= 2 and counters["todos"]:
from collections import Counter
for msg, cnt in Counter(counters["todos"]).most_common(): print(f' {cnt:3d}x {msg}')
atexit.register(_asm_gemm_report)
def can_use_asm_gemm(a:Tensor, b:Tensor) -> bool:
if a.dtype != b.dtype: return todo(f"dtypes must match {a.dtype} != {b.dtype}")
if a.dtype not in {dtypes.bfloat16, dtypes.float16, FP8_DTYPE}: return todo(f"only bfloat16/float16/fp8, got {a.dtype}")
batch, M, K = (1, *a.shape) if a.ndim == 2 else a.shape
N = b.shape[1]
if isinstance(a.device, tuple):
if a.ndim == 2 and a.uop.axis == 0 and b.uop.axis is None: M //= len(a.device)
elif a.ndim == 2 and a.uop.axis == 1 and b.uop.axis == 0: K //= len(a.device)
elif a.ndim == 2 and a.uop.axis is None and b.uop.axis == 1: N //= len(a.device)
elif a.ndim == 3 and a.uop.axis == 0 and b.uop.axis is None: batch //= len(a.device)
elif a.ndim == 3 and a.uop.axis == 1 and b.uop.axis is None: M //= len(a.device)
elif a.ndim == 3 and a.uop.axis is None and b.uop.axis == 1: N //= len(a.device)
elif a.ndim == 3 and a.uop.axis == 2 and b.uop.axis == 0: K //= len(a.device)
else: return todo(f"sharding mismatch a.ndim={a.ndim} a.uop.axis={a.uop.axis} b.uop.axis={b.uop.axis}")
dname = a.device[0]
else: dname = a.device
arch = Device[dname].renderer.target.arch
if batch not in {1, 2}: return todo(f"GEMM batch size {batch}")
if (M % TILE_M != 0 or N % TILE_N != 0 or K % TILE_K != 0) and arch == "gfx950":
return todo(f"GEMM shape ({M},{N},{K}) not a multiple of ({TILE_M},{TILE_N},{TILE_K})")
return True
# ** UOp gemm to test Tensor.custom_kernel multi and backward correctness on non cdna4
# note: this can be removed after we have GEMM on mixins
def custom_uop_gemm(C:UOp, A:UOp, B:UOp) -> UOp:
M, K = A.shape[0]*A.shape[1], A.shape[2]
K2, N = B.shape[(1 if B.ndim == 3 else 0):]
assert K == K2
m = UOp.range(M, 1)
n = UOp.range(N, 2)
k = UOp.range(K, 0, AxisType.REDUCE)
mul = (A.flatten().index((m*UOp.const(K)+k))*
B.flatten().index((k*UOp.const(N)+n))).cast(dtypes.float32)
red = mul.reduce(k, arg=Ops.ADD, dtype=dtypes.float32).cast(C.dtype)
store = C.flatten().index((m*UOp.const(N)+n)).store(red).end(m, n)
return store.sink(arg=KernelInfo(name=f'uop_gemm_{M}_{N}_{K}'))
# ** bf16 A @ B.T kernel in C
@functools.cache
def custom_hk_bf16_gemm(C:UOp, A:UOp, B:UOp, *args:UOp, dname:str) -> UOp:
M, K = A.shape[0]*A.shape[1], A.shape[2]
N, K2 = B.shape[(1 if B.ndim == 3 else 0):]
assert K == K2, f"{A.shape} {B.shape}"
block_m, block_n, block_k, num_warps = 256, 256, 64, 8
assert M % block_m == 0 and N % block_n == 0 and K % block_k == 0, f"invalid bf16 tile {(block_m, block_n, block_k)} for {(M, N, K)}"
threads = UOp.special(64 * num_warps, "lidx0")
workgroups = UOp.special((M // block_m) * (N // block_n), "gidx0")
b_extra = args[0].base if len(args) >= 1 else B.base
sink = UOp.sink(C.base, A.base, B.base, b_extra, threads, workgroups,
arg=KernelInfo(f"hk_bf16_gemm_{M}_{N}_{K}", estimates=Estimates(ops=2*M*N*K, mem=(M*K+N*K+M*N)*A.dtype.itemsize)))
kittens_path = pathlib.Path(__file__).parent.parent/"thunder"/"amd"
src = (kittens_path/"gemm_bf16.cpp").read_text()
lib = HIPCCCompiler("gfx950", [f"-I{(kittens_path/'include').as_posix()}", "-std=c++20", "-DKITTENS_CDNA4", "-ffast-math",
"-DHIP_ENABLE_WARP_SYNC_BUILTINS", f"-DGEMM_M={M}", f"-DGEMM_N={N}", f"-DGEMM_K={K}"]).compile_cached(src)
return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=src),
UOp(Ops.BINARY, arg=lib)))
@functools.cache
def custom_hk_bf16_atb_gemm(C:UOp, A:UOp, B:UOp, dname:str) -> UOp:
K, M = A.shape[0]*A.shape[1], A.shape[2]
K2, N = B.shape[0]*B.shape[1], B.shape[2]
assert K == K2, f"{A.shape} {B.shape}"
block_m, block_n, block_k, num_warps = 256, 256, 64, 8
assert M % block_m == 0 and N % block_n == 0 and K % block_k == 0, f"invalid bf16 atb tile {(block_m, block_n, block_k)} for {(M, N, K)}"
threads = UOp.special(64 * num_warps, "lidx0")
workgroups = UOp.special((M // block_m) * (N // block_n), "gidx0")
sink = UOp.sink(C.base, A.base, B.base, threads, workgroups,
arg=KernelInfo(f"hk_bf16_atb_gemm_{M}_{N}_{K}", estimates=Estimates(ops=2*M*N*K, mem=(M*K+N*K+M*N)*A.dtype.itemsize)))
kittens_path = pathlib.Path(__file__).parent.parent/"thunder"/"amd"
src = (kittens_path/"gemm_bf16_atb.cpp").read_text()
lib = HIPCCCompiler("gfx950", [f"-I{(kittens_path/'include').as_posix()}", "-std=c++20", "-DKITTENS_CDNA4", "-ffast-math",
"-DHIP_ENABLE_WARP_SYNC_BUILTINS", f"-DGEMM_M={M}", f"-DGEMM_N={N}", f"-DGEMM_K={K}"]).compile_cached(src)
return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=src),
UOp(Ops.BINARY, arg=lib)))
def hk_bf16_atb_gemm(a:Tensor, b:Tensor) -> Tensor:
assert a.dtype == b.dtype == dtypes.bfloat16, f"expected bf16, got {a.dtype} {b.dtype}"
assert a.ndim == b.ndim == 3 and a.shape[:2] == b.shape[:2], f"{a.shape} {b.shape}"
batch, rows, M = a.shape
N = b.shape[2]
assert M % TILE_M == 0 and N % TILE_N == 0 and (batch * rows) % TILE_K == 0, \
f"atb shape {a.shape} {b.shape} must produce (M,N,K) multiples of ({TILE_M},{TILE_N},{TILE_K})"
is_multi = isinstance(a.device, tuple)
reduce_out = False
if is_multi:
ndev = len(a.device)
if a.uop.axis in (0, 1) or b.uop.axis in (0, 1): inv, out_axis, reduce_out = Tensor.invalids(1, M, N, dtype=a.dtype, device=a.device), 0, True
elif b.uop.axis == 2: inv, out_axis = Tensor.invalids(1, M, N // ndev, dtype=a.dtype, device=a.device), 2
elif a.uop.axis == 2: inv, out_axis = Tensor.invalids(1, M // ndev, N, dtype=a.dtype, device=a.device), 1
else: inv, out_axis, reduce_out = Tensor.invalids(1, M, N, dtype=a.dtype, device=a.device), 0, True
out = Tensor(inv.uop.unshard(out_axis), device=a.device)
dname = a.device[0]
else:
out = Tensor.invalids(1, M, N, dtype=a.dtype, device=a.device)
dname = a.device
dname = dname.split(":")[0]
out = Tensor.custom_kernel(out, a, b, fxn=functools.partial(custom_hk_bf16_atb_gemm, dname=dname))[0]
if reduce_out: out = out.sum(0)
return out.squeeze(0) if out.ndim == 3 else out
# ** backward gemm, might use the asm gemm
def custom_gemm_bw(gradient:UOp, kernel:UOp, n_scales:int=2, has_grad_amax:bool=False, has_w_post:bool=False):
inputs = kernel.src[1:]
if inputs[1].dtype == FP8_DTYPE:
out, a, b = inputs[:3]
i = 3
s_x = inputs[i]; i += 1
has_w = n_scales >= 2
s_w = inputs[i] if has_w else None; i += has_w
s_g_amax = inputs[i] if n_scales == 3 else None; i += (n_scales == 3)
grad_amax_state = inputs[i] if has_grad_amax else None; i += has_grad_amax
next_grad_amax_state = inputs[i] if has_grad_amax else None; i += has_grad_amax
w_post = inputs[i] if has_w_post else None
a_t, b_t, g_t = Tensor(a, device=a.device), Tensor(b, device=a.device), Tensor(gradient, device=a.device)
s_x_t = Tensor(s_x, device=a.device)
s_w_t = Tensor(s_w, device=a.device) if has_w else None
s_g_amax_t = Tensor(s_g_amax, device=a.device) if s_g_amax is not None else None
w_post_t = Tensor(w_post, device=a.device) if has_w_post else None
g_t = g_t[:a.shape[0]]
from extra.llama_kernels.cast_amax import _grad_fp8_mailbox
from extra.llama_kernels.quantize_fp8_delayed import quantize_fp8_delayed
gbase = gradient.base if hasattr(gradient, "base") else gradient
mailbox_entry = _grad_fp8_mailbox.pop(gbase, None) or _grad_fp8_mailbox.pop(gradient, None)
if mailbox_entry is not None:
g_fp8_u, grad_amax_u = mailbox_entry
g_fp8 = Tensor(g_fp8_u, device=a.device)[:a.shape[0]]
g_amax = Tensor(grad_amax_u, device=a.device)
else:
assert grad_amax_state is not None, "fp8 matmul bwd needs either a mailbox entry or a grad_amax_state"
if getenv("CURRENT_GRAD_SCALE", 0):
g_fp8, _, g_amax = quantize_fp8(g_t, amax_state=None)
elif getenv("FUSED_GRAD_QUANTIZE", 0):
grad_amax_t = Tensor(grad_amax_state, device=a.device)
g_amax = grad_amax_t
g_fp8, _ = quantize_fp8_delayed(g_t, g_amax, Tensor(next_grad_amax_state, device=a.device))
else:
grad_amax_t = Tensor(grad_amax_state, device=a.device)
g_amax = grad_amax_t
g_fp8, _, new_grad_amax = quantize_fp8(g_t, amax_state=g_amax)
store_effect = next_grad_amax_state.store(new_grad_amax.uop)
g_fp8 = Tensor(g_fp8.contiguous().uop.after(store_effect), device=a.device)
# dgrad: applies grad/activation amax scales in the GEMM epilogue; w_scale is already inverse.
assert s_g_amax_t is None, "fp8 GEMM bwd through g_amax scaling is unsupported"
grad_a = asm_gemm(g_fp8, b_t, x_scale=s_x_t, w_scale=s_w_t, g_amax=g_amax) if has_w else asm_gemm(g_fp8, b_t, x_scale=s_x_t, g_amax=g_amax)
# wgrad: no w_scale
grad_b = hk_fp8_atb_gemm(g_fp8, a_t, x_scale=s_x_t, g_amax=g_amax)
# wgrad: rescale if not scalar
if w_post_t is not None:
grad_b = grad_b / w_post_t.reshape(*w_post_t.shape, *([1]*(grad_b.ndim - w_post_t.ndim)))
# one None per input: (out, a, b, x_scale[, w_scale][, grad_amax][, w_post_scale])
ret = (None, grad_a.uop, grad_b.uop) + tuple(None for _ in inputs[3:])
return ret
else:
hk_bf16 = len(inputs) == 4 and inputs[1].dtype == dtypes.bfloat16
if hk_bf16:
out, a, b_t, b = inputs
assert all_same([gradient.device, a.device, b_t.device, b.device, out.device])
else:
assert len(inputs) == 3, f"regular gemm must have exactly 3 sources, got: {len(inputs)}"
out, a, b = inputs
assert all_same([gradient.device, a.device, b.device, out.device])
a_t, b_t, g_t = Tensor(a, device=a.device), Tensor(b, device=a.device), Tensor(gradient, device=a.device)
g_t = g_t[:a.shape[0]]
if hk_bf16 and g_t.dtype != b_t.dtype: g_t = g_t.cast(b_t.dtype)
if can_use_asm_gemm(g_t, b_t.T): grad_a = asm_gemm(g_t, b_t.T).uop
else: grad_a = (g_t @ b_t.T).uop
if hk_bf16:
grad_b = hk_bf16_atb_gemm(a_t, g_t).uop
else:
a_t_flat, g_t_flat = a_t.permute(2, 0, 1).reshape(a_t.shape[2], -1), g_t.reshape(-1, g_t.shape[-1])
if can_use_asm_gemm(a_t_flat, g_t_flat): grad_b = asm_gemm(a_t_flat, g_t_flat).uop
else: grad_b = (a_t_flat @ g_t_flat).uop
# hk_bf16 uses b.T, writes gradients only for a and b
return (None, grad_a, None, grad_b) if hk_bf16 else (None, grad_a, grad_b)
# ** mxfp8 gemm backward
def custom_mx_gemm_bw(gradient:UOp, kernel:UOp, has_w_post:bool, w_stored:bool=False):
inputs = kernel.src[1:] # (out, a_q, b_q, a_si, b_si, a_e8, b_e8, [w_post])
aq, bq = Tensor(inputs[1], device=inputs[1].device), Tensor(inputs[2], device=inputs[2].device)
ae8, be8 = Tensor(inputs[5], device=inputs[5].device), Tensor(inputs[6], device=inputs[6].device)
wp = Tensor(inputs[7], device=inputs[7].device) if has_w_post else None
a_phys = (aq.reshape(-1, aq.shape[-1]).cast(dtypes.bfloat16) * _mx_block_scale(ae8)).cast(dtypes.bfloat16)
b_phys = (bq.cast(dtypes.bfloat16) * _mx_block_scale(be8)).cast(dtypes.bfloat16)
g = Tensor(gradient, device=aq.device)[:aq.shape[0]].reshape(aq.shape[0]*aq.shape[1], bq.shape[0]).cast(dtypes.bfloat16)
grad_a = asm_gemm(g, b_phys, mx=True)
grad_b = asm_gemm(g.T, a_phys, mx=True, a_pretranspose=g)
grad_a = (grad_a * _mx_block_scale(ae8)).reshape(aq.shape)
if not w_stored: grad_b = grad_b * _mx_block_scale(be8)
if wp is not None: grad_b = grad_b / wp.reshape(-1, 1)
return (None, grad_a.uop, grad_b.uop) + tuple(None for _ in inputs[3:])
# ** mxfp4 gemm backward
def custom_mxfp4_gemm_bw(gradient:UOp, kernel:UOp):
inputs = kernel.src[1:] # out, row operands/scales, BF16 operands, column operands/scales
assert len(inputs) == 11
a, w = Tensor(inputs[5], device=inputs[5].device), Tensor(inputs[6], device=inputs[6].device)
a_col, scale_a_col = Tensor(inputs[7], device=a.device), Tensor(inputs[8], device=a.device)
w_col, scale_w_col = Tensor(inputs[9], device=a.device), Tensor(inputs[10], device=a.device)
g = Tensor(gradient, device=a.device)[:a.shape[0]].cast(dtypes.bfloat16)
g_row, scale_g_row, g_col, scale_g_col = quantize_mxfp4(g, flatten_row=True)
grad_a = _mxfp4_gemm_quantized(g_row, w_col, scale_g_row, scale_w_col).reshape(*a.shape[:-1], w.shape[-1])
grad_w = _mxfp4_gemm_quantized(g_col, a_col, scale_g_col, scale_a_col).reshape(w.shape)
return (None, None, None, None, None, grad_a.uop, grad_w.uop, None, None, None, None)
# ** main gemm function
def asm_gemm(a:Tensor, b:Tensor, x_scale:Tensor|None=None, w_scale:Tensor|None=None, grad_amax_state:Tensor|None=None,
next_grad_amax_state:Tensor|None=None,
w_post_scale:Tensor|None=None, mx:bool=False, mx_scales:tuple|None=None, mx_w_stored:bool=False, g_amax:Tensor|None=None,
a_pretranspose:Tensor|None=None, mxfp4:bool=False) -> Tensor:
assert can_use_asm_gemm(a, b), f"{counters['todos'][-1]}"
if mxfp4:
assert not mx and mx_scales is None, "mxfp4 owns quantization; mx/mx_scales are for mxfp8"
assert a.dtype == dtypes.bfloat16, f"cannot quantize {a.dtype} to mxfp4"
counters["used"] += 1
unfold_batch = a.ndim == 3 and isinstance(a.device, tuple) and a.uop.axis == 2 and b.uop.axis == 0
if unfold_batch:
orig_batch = a.shape[0]
a = a.reshape(a.shape[0]*a.shape[1], a.shape[2])
squeeze = a.ndim == 2
if squeeze: a = a.unsqueeze(0)
out_dtype = dtypes.bfloat16 if a.dtype == FP8_DTYPE or mxfp4 else a.dtype
batch, M, K = a.shape
N = b.shape[1]
is_multi = isinstance(a.device, tuple)
if (k_sharded:=is_multi and a.uop.axis == 2): K //= len(a.device)
if (m_sharded:=is_multi and a.uop.axis == 1): M //= len(a.device)
n_sharded = is_multi and b.uop.axis == 1
if is_multi:
if n_sharded:
out = Tensor(Tensor.invalids(batch, M, N//len(a.device), dtype=out_dtype, device=a.device).uop.unshard(2), device=a.device)
elif m_sharded:
out = Tensor(Tensor.invalids(batch, M, N, dtype=out_dtype, device=a.device).uop.unshard(1), device=a.device)
else:
out = Tensor(Tensor.invalids(batch//len(a.device) if a.uop.axis==0 else batch, M, N, dtype=out_dtype, device=a.device).uop.unshard(0),
device=a.device)
else:
out = Tensor.invalids(batch, M, N, dtype=out_dtype, device=a.device)
renderer = Device[dname:=(a.device[0] if is_multi else a.device)].renderer
dname, arch = dname.split(":")[0], renderer.target.arch
if arch.startswith("gfx950") and getenv("USE_ASM", 1):
if mxfp4:
tile_m, tile_n = next((tm, tn) for tm, tn in ((256, 256), (192, 256), (128, 512)) if (batch*M) % tm == N % tn == 0)
fxn = functools.partial(custom_mxfp4_gemm, tile_m=tile_m, tile_n=tile_n)
w = b.T
a_q, scale_a, a_col, scale_a_col = quantize_mxfp4(a, shuffle_col=True)
b_q, scale_b, b_col, scale_b_col = quantize_mxfp4(w, shuffle_row=True, shuffle_col=True)
out = Tensor.custom_kernel(out, a_q, b_q, scale_a, scale_b, a, w,
a_col, scale_a_col, b_col, scale_b_col, fxn=fxn, grad_fxn=custom_mxfp4_gemm_bw)[0]
elif mx:
# mxfp8 1x32 block scaling
if mx_scales is not None:
a_si, a_e8, b_si, b_e8 = mx_scales
a_q, b_q = a.reshape(-1, a.shape[-1]), b.T
elif (a_pretranspose is not None and getenv("FUSED_GRAD_QUANTIZE", 0) and a_pretranspose.dtype == dtypes.bfloat16
and a_pretranspose.shape[0] % 32 == 0 and a_pretranspose.shape[1] % 256 == 0):
from extra.llama_kernels.transpose_quantize_mxfp8 import transpose_quantize_mxfp8
a_q, a_e8, a_si = transpose_quantize_mxfp8(a_pretranspose)
b_q, b_e8, b_si = quantize_mxfp8(b.T)
else:
a_q, a_e8, a_si = quantize_mxfp8(a.reshape(-1, a.shape[-1]))
b_q, b_e8, b_si = quantize_mxfp8(b.T)
has_w_post = w_post_scale is not None
fxn = functools.partial(custom_hk_mxfp8_gemm, dname=dname)
grad_fxn = functools.partial(custom_mx_gemm_bw, has_w_post=has_w_post, w_stored=mx_w_stored)
extra = [w_post_scale] if w_post_scale is not None else []
out = Tensor.custom_kernel(out, a_q.reshape(a.shape), b_q, a_si, b_si, a_e8, b_e8, *extra, fxn=fxn, grad_fxn=grad_fxn)[0]
# fp8 gemm computes a@b.T, kernel multiplies output by x_scale * w_scale before bf16 store
elif a.dtype == FP8_DTYPE:
scales = tuple(s for s in (x_scale, w_scale, g_amax) if s is not None)
scale_mode = (1 if x_scale is not None else 0) | (2 if w_scale is not None else 0) | (4 if g_amax is not None else 0)
assert (grad_amax_state is None) == (next_grad_amax_state is None)
extra = ([grad_amax_state, next_grad_amax_state] if grad_amax_state is not None else []) + ([w_post_scale] if w_post_scale is not None else [])
fxn = functools.partial(custom_hk_fp8_gemm, dname=dname, scale_mode=scale_mode)
bw = functools.partial(custom_gemm_bw, n_scales=len(scales), has_grad_amax=grad_amax_state is not None, has_w_post=w_post_scale is not None)
out = Tensor.custom_kernel(out, a, b.T, *scales, *extra, fxn=fxn, grad_fxn=bw)[0]
elif a.dtype == dtypes.bfloat16:
out = Tensor.custom_kernel(out, a, b.T, b, fxn=functools.partial(custom_hk_bf16_gemm, dname=dname), grad_fxn=custom_gemm_bw)[0]
else:
out = Tensor.custom_kernel(out, a, b, fxn=custom_uop_gemm, grad_fxn=custom_gemm_bw)[0]
if k_sharded: out = out.sum(0)
out = out.squeeze(0) if squeeze else out
if unfold_batch: out = out.reshape(orig_batch, -1, out.shape[-1])
if w_post_scale is not None: out = (out * w_post_scale.reshape(*([1]*(out.ndim-1)), -1)).cast(out.dtype)
return out

107
artifacts/package_sources/tinygrad/extra/gemm/cuda_matmul.py Normal file
View File

@@ -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)

44
artifacts/package_sources/tinygrad/extra/gemm/fuzz_matmul.py Normal file
View File

@@ -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")

194
artifacts/package_sources/tinygrad/extra/gemm/gemm.c Normal file
View File

@@ -0,0 +1,194 @@
// single: clang -O2 -march=native gemm.c
// multi: clang -O2 -march=native gemm.c -DNTHREADS=32 -lpthread
#define _GNU_SOURCE
// https://en.wikichip.org/wiki/amd/microarchitectures/zen_2
#include <stdint.h>
#include <time.h>
#include <sched.h>
#include <stdio.h>
#include <assert.h>
#include <math.h>
#include <string.h>
#include <immintrin.h>
#include <pthread.h>
#include <unistd.h>
#include <stdatomic.h>
//#define DEBUG
#ifdef DEBUG
#define N 8
#endif
#ifndef N
// NOTE: if you change this you have to rerun gemm.py
#define N 512
#endif
#ifndef NTHREADS
#define NTHREADS 1
#endif
// aligned?
float A[N*N] __attribute__ ((aligned (64)));
float B[N*N] __attribute__ ((aligned (64)));
float C[N*N] __attribute__ ((aligned (64)));
float val[N*N] __attribute__ ((aligned (64)));
__m256 *Am = (__m256*)A;
__m256 *Bm = (__m256*)B;
__m256 *Cm = (__m256*)C;
uint64_t nanos() {
struct timespec start;
clock_gettime(CLOCK_MONOTONIC_RAW, &start);
return (uint64_t)start.tv_sec*1000000000 + (uint64_t)start.tv_nsec;
}
float Bf[N*N] __attribute__ ((aligned (64)));
__m256 *Bfm = (__m256*)Bf;
#define BLOCK 8
#define BLOCK_Y 4
#define BLOCK_X 2
void matmul(int sy, int ey) {
// 136.77 GFLOPS on single core numpy
// 4.9 GHz is max boost for 5950X
// 32 FLOPS/cycle (16 FMAs, aka 2x 8 single wide / 32 byte FMAs)
// theoretical max is 156.8 GFLOPS, we see 150
// multicore theo max = 2508.8 GFLOPS, we see 1501.434299
// Bf = (y/8, k, 8)
for (int y = sy; y < ey; y+=BLOCK_Y) {
for (int x = 0; x < N; x+=BLOCK*BLOCK_X) {
__m256 acc[BLOCK_Y][BLOCK_X] = {};
for (int k = 0; k < N; k++) {
for (int iy = 0; iy < BLOCK_Y; iy++) {
__m256 ta = _mm256_broadcast_ss(&A[(y+iy)*N + k]);
for (int ix = 0; ix < BLOCK_X; ix++) {
acc[iy][ix] = _mm256_fmadd_ps(ta, Bfm[((x+ix*BLOCK)*N + k*8)/8], acc[iy][ix]);
}
}
}
for (int iy = 0; iy < BLOCK_Y; iy++) {
for (int ix = 0; ix < BLOCK_X; ix++) {
Cm[((y+iy)*N + x + ix * BLOCK)/8] = acc[iy][ix];
}
}
}
}
}
pthread_mutex_t lock = PTHREAD_MUTEX_INITIALIZER;
atomic_int nready = 0;
atomic_int ndone = 0;
void *matmul_thread(void *n) {
int k = (int)(int64_t)n;
int sy = (N/NTHREADS) * k;
int ey = (N/NTHREADS) * (k+1);
cpu_set_t set;
CPU_ZERO(&set);
CPU_SET(k,&set);
pthread_setaffinity_np(pthread_self(), sizeof(cpu_set_t), &set);
nready++;
// gotta have main lock once to signal start
pthread_mutex_lock(&lock);
pthread_mutex_unlock(&lock);
matmul(sy, ey);
// we done
ndone++;
return NULL;
}
int main() {
printf("hello with %d threads\n", NTHREADS);
#ifdef DEBUG
for (int i = 0; i < N*N; i++) A[i] = i;
for (int i = 0; i < N*N; i++) B[i] = i;
#else
FILE *f = fopen("/tmp/matmul", "rb");
if (f == NULL) {
printf("please pregenerate python /tmp/matmul file\n");
return -1;
}
fread(A, 1, sizeof(float)*N*N, f);
fread(B, 1, sizeof(float)*N*N, f);
fread(val, 1, sizeof(float)*N*N, f);
fclose(f);
#endif
// preswizzle
for (int y = 0; y < N; y+=8) {
for (int x = 0; x < N; x++) {
for (int iy = 0; iy < 8; iy++) {
Bf[y*N + x*8 + iy] = B[(y+iy)*N + x];
}
}
}
for (int i = 0; i < 10; i++) {
memset(C, 0, N*N*sizeof(float));
#if NTHREADS != 1
nready = 0;
ndone = 0;
pthread_mutex_lock(&lock);
pthread_t threads[NTHREADS];
for (int j = 0; j < NTHREADS; j++) {
pthread_create(&threads[j], NULL, matmul_thread, (void *)(uint64_t)j);
}
while (nready != NTHREADS) usleep(1);
#endif
uint64_t start = nanos();
#if NTHREADS == 1
matmul(0, N);
#else
// unlocking mutex starts threads
pthread_mutex_unlock(&lock);
while (ndone != NTHREADS) usleep(1);
#endif
uint64_t end = nanos();
#if NTHREADS != 1
for (int j = 0; j < NTHREADS; j++) {
pthread_join(threads[j], NULL);
}
#endif
double gflop = (2.0*N*N*N)*1e-9;
double s = (end-start)*1e-9;
printf("%f GFLOP/S -- %.2f ms\n", gflop/s, s*1e3);
// hack around throttling
//if (i%4 == 0) sleep(1);
}
#ifdef DEBUG
for (int i = 0; i < N*N; i++) {
if (i%N == 0 && i != 0) printf("\n");
printf("%f ", C[i]);
}
printf("\n");
#else
for (int k = 0; k < N*N; k++) {
if (fabsf(C[k] - val[k]) > 1e-3) {
printf("MISMATCH AT %d, %f != %f\n", k, C[k], val[k]);
return -1;
}
}
printf("match\n");
#endif
return 0;
}

28
artifacts/package_sources/tinygrad/extra/gemm/gemm.py Executable file
View File

@@ -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)

4124
artifacts/package_sources/tinygrad/extra/gemm/gemm_mxfp4.py Normal file
View File

File diff suppressed because it is too large Load Diff

90
artifacts/package_sources/tinygrad/extra/gemm/halide_gemm.py Normal file
View File

@@ -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.")

142
artifacts/package_sources/tinygrad/extra/gemm/hip_matmul.py Normal file
View File

@@ -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)

508
artifacts/package_sources/tinygrad/extra/gemm/max_kernels/nv.fp16_fp16_fp16.2_stage.cu Normal file
View File

@@ -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;
}

465
artifacts/package_sources/tinygrad/extra/gemm/max_kernels/nv.fp16_fp16_fp16.3_stage.cu Normal file
View File

@@ -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;
}

517
artifacts/package_sources/tinygrad/extra/gemm/max_kernels/nv.fp16_fp16_fp16.3_stage_swizzled.cu Normal file
View File

@@ -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;
}

482
artifacts/package_sources/tinygrad/extra/gemm/max_kernels/nv.fp16_fp16_fp16.max.cu Normal file
View File

@@ -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();
}

486
artifacts/package_sources/tinygrad/extra/gemm/max_kernels/nv.fp16_fp16_fp16.no_xor.cu Normal file
View File

@@ -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();
}

157
artifacts/package_sources/tinygrad/extra/gemm/max_kernels/nv.fp16_fp32_fp16.hcopt.cu Normal file
View File

@@ -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));
}

398
artifacts/package_sources/tinygrad/extra/gemm/max_kernels/nv.fp16_fp32_fp32.2_stage_swizzled_smem_input.cu Normal file
View File

@@ -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;
}

363
artifacts/package_sources/tinygrad/extra/gemm/max_kernels/nv.fp16_fp32_fp32.flat_smem_input.cu Normal file
View File

@@ -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;
}

439
artifacts/package_sources/tinygrad/extra/gemm/max_kernels/nv.fp16_fp32_fp32.max.cu Normal file
View File

@@ -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;
}

371
artifacts/package_sources/tinygrad/extra/gemm/max_kernels/nv.fp16_fp32_fp32.swizzled_smem_input.cu Normal file
View File

@@ -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;
}

218
artifacts/package_sources/tinygrad/extra/gemm/max_matmul.py Normal file
View File

@@ -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")

49
artifacts/package_sources/tinygrad/extra/gemm/metal_conv.py Normal file
View File

@@ -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")

132
artifacts/package_sources/tinygrad/extra/gemm/metal_matmul.py Normal file
View File

@@ -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")

113
artifacts/package_sources/tinygrad/extra/gemm/metal_matvec.py Normal file
View File

@@ -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)

39
artifacts/package_sources/tinygrad/extra/gemm/metal_uop_matmul.py Normal file
View File

@@ -0,0 +1,39 @@
from tinygrad import UOp, dtypes
from tinygrad.uop.ops import AxisType, 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.stack(*[mat_idx(a, gx, gk, warp, i) for i in range(2)])
b_tc = UOp.stack(*[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)
acc_load = UOp.stack(acc.after(gk)[0], acc.after(gk)[1])
out = UOp.wmma(a_tc, b_tc, acc_load, (8, 8, 8), 'METAL', 32)
end_loop = UOp.group(*[acc[i].store(out.index(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)

227
artifacts/package_sources/tinygrad/extra/gemm/mi350x_uop_matmul.py Normal file
View File

@@ -0,0 +1,227 @@
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
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, 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((0.0,)*4, dtypes.float), 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, slot=1, addrspace=AddrSpace.REG)
Br = UOp.placeholder((BLOCK_N//TC_N,), dtypes.half, 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
out = UOp.wmma(Ar[M_inner_loop], Br[N_inner_loop], acc_load, (16, 16, 32), 'AMD', 64)
# 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].index(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.stack(*[acc.after(K_loop)[i] for i in range(4)])
out = UOp.wmma(A_in, B_in, acc_load, (16, 16, 32), 'AMD', 64)
# store back the acc
acc = acc.after(UOp.group(*[acc[i].store(out.index(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"

140
artifacts/package_sources/tinygrad/extra/gemm/mi350x_uop_matmul_2.py Normal file
View File

@@ -0,0 +1,140 @@
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
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.WEAK) 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, slot=1, addrspace=AddrSpace.REG)
Br = UOp.placeholder((BLOCK_N//TC_N,), dtypes.half, 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
out = UOp.wmma(Ar[M_inner_loop], Br[N_inner_loop], acc_load, (16, 16, 32), 'AMD', 64)
# 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).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, 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"

111
artifacts/package_sources/tinygrad/extra/gemm/moe_gemm.py Normal file
View File

@@ -0,0 +1,111 @@
import functools, pathlib
from tinygrad import Tensor, dtypes
from tinygrad.uop.ops import UOp, Ops, KernelInfo
from tinygrad.renderer import Estimates
from tinygrad.runtime.support.compiler_amd import HIPCCCompiler
from extra.gemm.cdna_asm_gemm import quantize_mxfp8, _mx_block_scale, _mx_block_scale_3d
@functools.cache
def custom_hk_grouped_mxfp8_gemm(C:UOp, A:UOp, B:UOp, scale_A:UOp, scale_B:UOp, *extra:UOp, dname:str, n_experts:int) -> UOp:
M, K = A.shape
E, N, K2 = B.shape
assert K == K2, f"{A.shape} {B.shape}"
assert E == n_experts, f"{E} != {n_experts}"
threads = UOp.special(64 * 8, "lidx0")
workgroups = UOp.special((M // 256) * (N // 256), "gidx0")
sink_inputs = (C.base, A.base, B.base, scale_A.base, scale_B.base, extra[0].base, extra[1].base, extra[2].base, threads, workgroups)
sink = UOp.sink(*sink_inputs,
arg=KernelInfo(f"hk_grouped_mxfp8_gemm_{E}_{M}_{N}_{K}",
estimates=Estimates(ops=2*M*N*K, mem=(M*K+E*N*K)*A.dtype.itemsize+M*N*C.dtype.itemsize)))
kittens_path = pathlib.Path(__file__).parent.parent/"thunder"/"amd"
src = (kittens_path/"grouped_mxfp8_gemm.cpp").read_text()
lib = HIPCCCompiler("gfx950", [f"-I{(kittens_path/'include').as_posix()}", "-std=c++20", "-DKITTENS_CDNA4", "-ffast-math",
"-DHIP_ENABLE_WARP_SYNC_BUILTINS", f"-DGEMM_M={M}", f"-DGEMM_N={N}", f"-DGEMM_K={K}",
f"-DGEMM_E={E}"]).compile_cached(src)
return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=src),
UOp(Ops.BINARY, arg=lib)))
@functools.cache
def custom_hk_grouped_mxfp8_wgrad(C:UOp, A:UOp, B:UOp, scale_A:UOp, scale_B:UOp, expert_off:UOp, *, dname:str, n_experts:int) -> UOp:
N, M = A.shape
K, M2 = B.shape
assert M == M2, f"{A.shape} {B.shape}"
E = n_experts
threads = UOp.special(64 * 8, "lidx0")
workgroups = UOp.special(E * (N // 256) * (K // 256), "gidx0")
sink = UOp.sink(C.base, A.base, B.base, scale_A.base, scale_B.base, expert_off.base, threads, workgroups,
arg=KernelInfo(f"hk_grouped_mxfp8_wgrad_{E}_{M}_{N}_{K}",
estimates=Estimates(ops=2*M*N*K, mem=(N*M+K*M)*A.dtype.itemsize+E*N*K*C.dtype.itemsize)))
kittens_path = pathlib.Path(__file__).parent.parent/"thunder"/"amd"
src = (kittens_path/"grouped_mxfp8_wgrad.cpp").read_text()
lib = HIPCCCompiler("gfx950", [f"-I{(kittens_path/'include').as_posix()}", "-std=c++20", "-DKITTENS_CDNA4", "-ffast-math",
"-DHIP_ENABLE_WARP_SYNC_BUILTINS", f"-DWGRAD_M={M}", f"-DWGRAD_N={N}", f"-DWGRAD_K={K}",
f"-DWGRAD_E={E}"]).compile_cached(src)
return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=src),
UOp(Ops.BINARY, arg=lib)))
def grouped_mx_wgrad(g:Tensor, xg:Tensor, expert_off:Tensor, n_experts:int) -> Tensor:
from extra.llama_kernels.transpose_quantize_mxfp8 import transpose_quantize_mxfp8
M, N = g.shape
M2, K = xg.shape
assert M == M2, f"{g.shape} {xg.shape}"
assert M % 128 == 0 and N % 256 == 0 and K % 256 == 0, f"wgrad needs M%128,N%256,K%256, got {g.shape} {xg.shape}"
gT, _, g_si = transpose_quantize_mxfp8(g.contiguous())
xT, _, x_si = transpose_quantize_mxfp8(xg.contiguous())
dname = (g.device[0] if isinstance(g.device, tuple) else g.device).split(":")[0]
is_multi = isinstance(g.device, tuple)
inv = Tensor.invalids(1, n_experts * N, K, dtype=dtypes.bfloat16, device=g.device)
out = Tensor(inv.uop.unshard(0), device=g.device) if is_multi else inv
out = Tensor.custom_kernel(out, gT, xT, g_si, x_si, expert_off,
fxn=functools.partial(custom_hk_grouped_mxfp8_wgrad, dname=dname, n_experts=n_experts))[0]
out = out.sum(0) if is_multi else out.squeeze(0)
return out.reshape(n_experts, N, K)
def mx_pack_3d(e8:Tensor) -> Tensor:
E, rows, scale_K = e8.shape
return e8.reshape(E, rows, scale_K // 4, 4).bitcast(dtypes.uint32).reshape(E, rows, scale_K // 4).permute(0, 2, 1).contiguous()
@functools.cache
def custom_grouped_mx_gemm_bw(gradient:UOp, kernel:UOp, w_stored:bool=False) -> tuple:
inputs = kernel.src[1:]
aq = Tensor(inputs[1], device=inputs[1].device)
bq = Tensor(inputs[2], device=inputs[2].device)
ae8 = Tensor(inputs[5], device=inputs[5].device)
be8 = Tensor(inputs[6], device=inputs[6].device)
E, N = bq.shape[0], bq.shape[1]
M, K = aq.shape
g = Tensor(gradient, device=aq.device).reshape(M, N).cast(dtypes.bfloat16)
x_phys = (aq.cast(dtypes.bfloat16) * _mx_block_scale(ae8).cast(dtypes.bfloat16))
w_phys = (bq.cast(dtypes.bfloat16) * _mx_block_scale_3d(be8).cast(dtypes.bfloat16))
expert_off = Tensor(inputs[7], device=inputs[7].device)
grad_x = grouped_mx_gemm(g, w_phys.transpose(1, 2), expert_off)
grad_w = grouped_mx_wgrad(g, x_phys, expert_off, E)
grad_xq = grad_x * _mx_block_scale(ae8).cast(dtypes.bfloat16)
grad_wq = grad_w.contiguous() if w_stored else (grad_w * _mx_block_scale_3d(be8).cast(dtypes.bfloat16)).contiguous()
return (None, grad_xq.uop, grad_wq.uop) + tuple(None for _ in inputs[3:])
_grouped_bw_stored = functools.partial(custom_grouped_mx_gemm_bw, w_stored=True)
def grouped_mx_gemm(x:Tensor, w:Tensor|tuple[Tensor, Tensor], expert_off:Tensor) -> Tensor:
if (pre_quantized := isinstance(w, tuple)):
w_q, w_e8 = w
E, N, K2 = w_q.shape
else:
E, N, K2 = w.shape
M, K = x.shape
assert K == K2, f"shape mismatch {x.shape} {w.shape}"
assert M % 256 == 0 and N % 256 == 0 and K % 128 == 0, f"grouped mxfp8 needs M%256,N%256,K%128, got {x.shape} {w.shape}"
dname = (x.device[0] if isinstance(x.device, tuple) else x.device).split(":")[0]
x_q, x_e8, x_si = quantize_mxfp8(x)
if not pre_quantized: w_q, w_e8, _ = quantize_mxfp8(w)
w_si = mx_pack_3d(w_e8)
xe_in, out_shape = x_e8.reshape(M, K // 32), (M, N)
if isinstance(x.device, tuple) and (row_axis := x.uop.axis) is not None:
ndev = len(x.device)
out = Tensor(Tensor.invalids(*(s // ndev if i == row_axis else s for i, s in enumerate(out_shape)),
dtype=dtypes.bfloat16, device=x.device).uop.unshard(row_axis), device=x.device)
else:
out = Tensor.invalids(*out_shape, dtype=dtypes.bfloat16, device=x.device)
return Tensor.custom_kernel(out, x_q, w_q, x_si, w_si, xe_in, w_e8, expert_off,
fxn=functools.partial(custom_hk_grouped_mxfp8_gemm, dname=dname, n_experts=E),
grad_fxn=(_grouped_bw_stored if pre_quantized else custom_grouped_mx_gemm_bw))[0]

130
artifacts/package_sources/tinygrad/extra/gemm/moe_routing.py Normal file
View File

@@ -0,0 +1,130 @@
from tinygrad import Tensor, dtypes
from tinygrad.uop.ops import UOp, Ops, KernelInfo, AxisType
BLOCK_ROW = 256
def _sharded_invalids(shape:tuple[int, ...], dtype, device) -> Tensor:
if isinstance(device, tuple):
return Tensor.invalids(*shape, dtype=dtype, device=device[0]).shard(device, axis=0)
return Tensor.invalids(*shape, dtype=dtype, device=device)
def _atomic_add(device:str) -> str:
return "__hip_atomic_fetch_add({0}, {1}, __ATOMIC_RELAXED, __HIP_MEMORY_SCOPE_AGENT);" if device == "AMD" \
else "__atomic_fetch_add({0}, {1}, __ATOMIC_RELAXED);"
def _blk_for(D:int) -> int:
blk = 64
while D % blk: blk //= 2
return blk
def _kv_ranges(G, N, D, BLK):
g = UOp.range(G, 0)
m = UOp.range(N, 1)
jo = UOp.range(D // BLK, 2)
ji = UOp.range(BLK, 3, AxisType.LOCAL)
return g, m, jo * BLK + ji, jo, ji
def _ggather_fwd_kernel(out:UOp, table:UOp, idx:UOp) -> UOp:
G, M, D = out.shape
g, m, j, jo, ji = _kv_ranges(G, M, D, _blk_for(D))
row = idx.index(g, m).cast(dtypes.weakint)
val = table.index(g, row, j).load()
return out.index(g, m, j).store(val).end(g, m, jo, ji).sink(
arg=KernelInfo(name=f"ggather_fwd_{M}_{D}", opts_to_apply=()))
def _ggather_zero_kernel(out:UOp) -> UOp:
i = UOp.range(out.numel(), 0)
return out.flatten().index(i).store(UOp.const(0.0, out.dtype)).end(i).sink(arg=KernelInfo(name="ggather_zero"))
def _sharded_zeros(shape:tuple[int, ...], dtype, device) -> Tensor:
return Tensor.custom_kernel(_sharded_invalids(shape, dtype, device), fxn=_ggather_zero_kernel)[0]
def _ggather_bwd(gradient:UOp, kernel:UOp) -> tuple:
_, table_u, idx_u = kernel.src[1:4]
dev = table_u.device
device = (dev[0] if isinstance(dev, tuple) else dev).split(":")[0]
G, R, D = table_u.shape
gt = _sharded_zeros((G, R, D), dtypes.float32, dev)
go = Tensor(gradient, device=dev)
atomic_str = _atomic_add(device)
def _bwd_kernel(gtab:UOp, gout:UOp, idx:UOp) -> UOp:
Gk, M, Dk = gout.shape
g, m, j, jo, ji = _kv_ranges(Gk, M, Dk, _blk_for(Dk))
row = idx.index(g, m).cast(dtypes.weakint)
val = gout.index(g, m, j).load().cast(dtypes.float32)
atomic = UOp(Ops.CUSTOM, dtypes.void, (gtab.index(g, row, j), val), arg=atomic_str)
return atomic.end(g, m, jo, ji).sink(arg=KernelInfo(name=f"ggather_bwd_{M}_{Dk}", opts_to_apply=()))
grad_table = Tensor.custom_kernel(gt, go, Tensor(idx_u, device=dev), fxn=_bwd_kernel)[0]
return (None, grad_table.cast(table_u.dtype).uop, None)
def grouped_gather_rows(table:Tensor, idx:Tensor, n_groups:int) -> Tensor:
G, R, D = table.shape
M = idx.shape[1]
out = _sharded_invalids((G, M, D), table.dtype, table.device)
return Tensor.custom_kernel(out, table, idx, fxn=_ggather_fwd_kernel, grad_fxn=_ggather_bwd)[0]
def _gscatter_fwd_kernel(out:UOp, src:UOp, idx:UOp) -> UOp:
G, M, D = out.shape
k = idx.shape[1] // src.shape[1]
g, m, j, jo, ji = _kv_ranges(G, idx.shape[1], D, _blk_for(D))
row = idx.index(g, m).cast(dtypes.weakint)
val = src.index(g, (m // k).cast(dtypes.weakint), j).load()
return out.index(g, row, j).store(val).end(g, m, jo, ji).sink(
arg=KernelInfo(name=f"gscatter_fwd_{idx.shape[1]}_{D}", opts_to_apply=()))
def _gscatter_bwd(gradient:UOp, kernel:UOp) -> tuple:
_, src_u, idx_u = kernel.src[1:4]
dev = src_u.device
G, T_l, D = src_u.shape
k = idx_u.shape[1] // T_l
sel = grouped_gather_rows(Tensor(gradient, device=dev), Tensor(idx_u, device=dev), G)
return (None, sel.reshape(G, T_l, k, D).sum(2).cast(src_u.dtype).uop, None)
def grouped_scatter_rows(src:Tensor, idx:Tensor, m_l:int) -> Tensor:
G, T_l, D = src.shape
zero = _sharded_zeros((G, m_l, D), src.dtype, src.device)
return Tensor.custom_kernel(zero, src, idx, fxn=_gscatter_fwd_kernel, grad_fxn=_gscatter_bwd)[0]
def m_max_for(t_local:int, experts_per_tok:int, n_experts:int) -> int:
return (-(-t_local * experts_per_tok // BLOCK_ROW) + n_experts) * BLOCK_ROW
class Routing:
def __init__(self, weights:Tensor, dest_row:Tensor, off:Tensor, m_l:int, n_groups:int, t_local:int):
self.weights, self.dest_row = weights, dest_row
self.off = off
self.m_l, self.n_groups, self.t_local = m_l, n_groups, t_local
@property
def rows_e(self) -> Tensor:
G, E = self.off.shape[0], self.off.shape[1] - 1
tr = Tensor.arange(self.m_l // BLOCK_ROW, dtype=dtypes.int32).reshape(1, -1, 1) * BLOCK_ROW
tr = tr.shard(self.off.device) if isinstance(self.off.device, tuple) else tr.to(self.off.device)
tile_e = ((tr >= self.off[:, :E].reshape(G, 1, E)).sum(-1) - 1).cast(dtypes.int32)
return tile_e.reshape(-1, 1).expand(-1, BLOCK_ROW).reshape(-1)
def n_groups_of(t:Tensor) -> int:
return len(t.device) if isinstance(t.device, tuple) else 1
def route(logits:Tensor, experts_per_tok:int, n_experts:int) -> Routing:
T, E = logits.shape
k, G = experts_per_tok, n_groups_of(logits)
assert T % G == 0, f"tokens {T} must split across {G} devices"
T_l, m_l = T // G, m_max_for(T // G, k, n_experts)
topv, topi = logits.reshape(G, T_l, E).topk(k)
weights = topv.softmax(-1)
m = topi.reshape(G, T_l * k).cast(dtypes.int32).one_hot(E).cast(dtypes.int32)
pad = ((m.sum(1) + (BLOCK_ROW - 1)) // BLOCK_ROW) * BLOCK_ROW
off = pad.cumsum(1).pad(((0, 0), (1, 0)))
dest_row = ((m.cumsum(1) + off[:, :E].reshape(G, 1, E)) * m).sum(-1).sub(1).cast(dtypes.int32)
return Routing(weights, dest_row, off, m_l, G, T_l)
def dispatch(x:Tensor, r:Routing) -> Tensor:
G, D = r.n_groups, x.shape[-1]
return grouped_scatter_rows(x.reshape(G, r.t_local, D), r.dest_row, r.m_l).reshape(G * r.m_l, D)
def combine(y:Tensor, r:Routing, n_tokens:int, experts_per_tok:int) -> Tensor:
G, D, k = r.n_groups, y.shape[-1], experts_per_tok
sel = grouped_gather_rows(y.reshape(G, r.m_l, D), r.dest_row, G).reshape(G, r.t_local, k, D)
return (sel * r.weights.reshape(G, r.t_local, k, 1).cast(sel.dtype)).sum(2).reshape(n_tokens, D).cast(y.dtype)

249
artifacts/package_sources/tinygrad/extra/gemm/rdna4_asm_matmul.py Normal file
View File

@@ -0,0 +1,249 @@
# 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_size = max(LDS_SIZE, 65536//getenv("LIMIT_OCC",2))
lds = UOp.placeholder((lds_size,), dtypes.uint8, 0, AddrSpace.LOCAL)
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.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()

20
artifacts/package_sources/tinygrad/extra/gemm/real_pmatmul.py Normal file
View File

@@ -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")

33
artifacts/package_sources/tinygrad/extra/gemm/simple_conv.py Normal file
View File

@@ -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)

57
artifacts/package_sources/tinygrad/extra/gemm/simple_matmul.py Normal file
View File

@@ -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

30
artifacts/package_sources/tinygrad/extra/gemm/simple_matvec.py Normal file
View File

@@ -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)

34
artifacts/package_sources/tinygrad/extra/gemm/tinygrad_nv_matmul.py Normal file
View File

@@ -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)

Some files were not shown because too many files have changed in this diff Show More