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

IQ.Pilot Release Commit @ b6534c0

Browse Source
This commit is contained in:
IQ.Lvbs CI [bot]
2026-08-27 20:17:33 -05:00
commit 00f07cac48
4706 changed files with 1257146 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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)

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

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

117
artifacts/package_sources/tinygrad/extra/gemm/triton_nv_matmul.py Normal file
View File

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

46
artifacts/package_sources/tinygrad/extra/gemm/tvm_gemm.py Normal file
View File

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