mirror of
https://github.com/tinygrad/tinygrad.git
synced 2026-08-29 12:56:07 +00:00
cdna asm_gemm in one file, remove old rdna3 asm (#15281)
This commit is contained in:
qazal
@@ -1,123 +0,0 @@
|
||||
import atexit, functools
|
||||
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
|
||||
from extra.gemm.asm.cdna.asm import build_kernel, TILE_M, TILE_N, TILE_K, NUM_WG
|
||||
|
||||
# ** CDNA4 assembly gemm
|
||||
|
||||
WORKGROUP_SIZE = 256
|
||||
|
||||
@functools.cache
|
||||
def custom_asm_gemm(C:UOp, A:UOp, B:UOp, dname:str) -> UOp:
|
||||
batch, M, K = A.shape
|
||||
K2, N = B.shape[(1 if B.ndim == 3 else 0):]
|
||||
assert K == K2
|
||||
lidx = UOp.special(WORKGROUP_SIZE, "lidx0")
|
||||
gidx = UOp.special(NUM_WG, "gidx0")
|
||||
insts = build_kernel(batch, M, N, K, A.dtype.base)
|
||||
lds = UOp(Ops.DEFINE_LOCAL, dtypes.uint8.ptr(size=133_120, addrspace=AddrSpace.LOCAL), (), 'lds')
|
||||
sink = UOp.sink(C.base, A.base, B.base, lds, lidx, gidx,
|
||||
arg=KernelInfo(name=f"gemm_{batch}_{M}_{N}_{K}", estimates=Estimates(ops=2*batch*M*N*K, mem=(batch*M*K + K*N + batch*M*N)*2)))
|
||||
return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=dname),
|
||||
UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in insts]))))
|
||||
|
||||
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}: return todo(f"only bfloat16/float16, 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 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 = getattr(Device[dname].renderer, "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, AxisType.LOOP)
|
||||
n = UOp.range(N, 2, AxisType.LOOP)
|
||||
k = UOp.range(K, 0, AxisType.REDUCE)
|
||||
mul = (A.flatten().index((m*UOp.const(dtypes.index, K)+k))*
|
||||
B.flatten().index((k*UOp.const(dtypes.index, N)+n))).cast(dtypes.float32)
|
||||
red = mul.reduce(k, arg=Ops.ADD, dtype=dtypes.float32).cast(C.dtype.base)
|
||||
store = C.flatten().index((m*UOp.const(dtypes.index, N)+n), ptr=True).store(red).end(m, n)
|
||||
return store.sink(arg=KernelInfo(name=f'uop_gemm_{M}_{N}_{K}'))
|
||||
|
||||
# ** backward gemm, might use the asm gemm
|
||||
|
||||
def custom_gemm_bw(gradient:UOp, kernel:UOp):
|
||||
out, a, b = kernel.src[1:]
|
||||
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)
|
||||
# TODO: this needs to be cleaned up and done properly, the batch dim of grad and a multi need to align
|
||||
g_t = g_t[:a.shape[0]]
|
||||
grad_a = (g_t @ b_t.T).uop
|
||||
grad_b = (a_t.permute(2, 0, 1).reshape(a_t.shape[2], -1) @ g_t.reshape(-1, g_t.shape[-1])).uop
|
||||
return (None, grad_a, grad_b)
|
||||
|
||||
# ** main gemm function
|
||||
|
||||
def asm_gemm(a:Tensor, b:Tensor) -> Tensor:
|
||||
assert can_use_asm_gemm(a, b), f"{counters['todos'][-1]}"
|
||||
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)
|
||||
|
||||
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.empty(batch, M, N//len(a.device), dtype=a.dtype, device=a.device).uop.multi(2), device=a.device)
|
||||
elif m_sharded:
|
||||
out = Tensor(Tensor.empty(batch, M, N, dtype=a.dtype, device=a.device).uop.multi(1), device=a.device)
|
||||
else:
|
||||
out = Tensor(Tensor.empty(batch//len(a.device) if a.uop.axis==0 else batch, M, N, dtype=a.dtype, device=a.device).uop.multi(0), device=a.device)
|
||||
else:
|
||||
out = Tensor.empty(batch, M, N, dtype=a.dtype, device=a.device)
|
||||
|
||||
renderer = Device[a.device[0] if is_multi else a.device].renderer
|
||||
dname, arch = renderer.device, getattr(renderer, "arch", "")
|
||||
if arch.startswith("gfx950") and getenv("USE_ASM", 1):
|
||||
out = Tensor.custom_kernel(out, a, b, fxn=functools.partial(custom_asm_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])
|
||||
return out
|
||||
File diff suppressed because it is too large
Load Diff
@@ -1,76 +0,0 @@
|
||||
.text
|
||||
.section .text.
|
||||
.global gemm
|
||||
.p2align 8
|
||||
.type gemm,@function
|
||||
|
||||
gemm:
|
||||
INSTRUCTIONS
|
||||
|
||||
.section .rodata,"a",@progbits
|
||||
.p2align 6, 0x0
|
||||
.amdhsa_kernel gemm
|
||||
# basic memory requirements
|
||||
.amdhsa_group_segment_fixed_size 30336
|
||||
.amdhsa_private_segment_fixed_size 0
|
||||
.amdhsa_kernarg_size 32
|
||||
# register usage (RSRC1)
|
||||
.amdhsa_next_free_vgpr 256
|
||||
.amdhsa_next_free_sgpr 100
|
||||
# workgroup / workitem IDs (RSRC2)
|
||||
.amdhsa_system_sgpr_workgroup_id_x 1
|
||||
.amdhsa_system_sgpr_workgroup_id_y 1
|
||||
.amdhsa_system_sgpr_workgroup_id_z 1
|
||||
# user SGPRs: kernarg ptr in s[0:1]
|
||||
.amdhsa_user_sgpr_kernarg_segment_ptr 1
|
||||
.amdhsa_user_sgpr_count 2
|
||||
# gfx10+ / gfx11 specifics (RSRC1[29..31])
|
||||
.amdhsa_wavefront_size32 1
|
||||
.amdhsa_workgroup_processor_mode 1
|
||||
.amdhsa_memory_ordered 1
|
||||
.amdhsa_forward_progress 1
|
||||
# misc for gfx11
|
||||
.amdhsa_dx10_clamp 1
|
||||
.amdhsa_ieee_mode 1
|
||||
.amdhsa_uses_dynamic_stack 0
|
||||
.end_amdhsa_kernel
|
||||
|
||||
.amdgpu_metadata
|
||||
---
|
||||
amdhsa.kernels:
|
||||
- .args:
|
||||
- .address_space: generic
|
||||
.name: C
|
||||
.offset: 0
|
||||
.size: 8
|
||||
.value_kind: global_buffer
|
||||
.value_type: f16
|
||||
- .address_space: generic
|
||||
.name: A
|
||||
.offset: 8
|
||||
.size: 8
|
||||
.value_kind: global_buffer
|
||||
.value_type: f16
|
||||
- .address_space: generic
|
||||
.name: B
|
||||
.offset: 16
|
||||
.size: 8
|
||||
.value_kind: global_buffer
|
||||
.value_type: f16
|
||||
.group_segment_fixed_size: 30336
|
||||
.kernarg_segment_align: 8
|
||||
.kernarg_segment_size: 32
|
||||
.max_flat_workgroup_size: 128
|
||||
.name: gemm
|
||||
.private_segment_fixed_size: 0
|
||||
.sgpr_count: 70
|
||||
.sgpr_spill_count: 0
|
||||
.symbol: gemm.kd
|
||||
.vgpr_count: 256
|
||||
.vgpr_spill_count: 0
|
||||
.wavefront_size: 32
|
||||
amdhsa.version:
|
||||
- 1
|
||||
- 1
|
||||
...
|
||||
.end_amdgpu_metadata
|
||||
@@ -1,30 +0,0 @@
|
||||
import math, pathlib
|
||||
|
||||
from tinygrad import Device, dtypes
|
||||
from tinygrad.uop.ops import UOp, Ops, KernelInfo
|
||||
|
||||
from extra.gemm.amd_uop_matmul import test_matmul
|
||||
|
||||
N = 4096
|
||||
TN = 96
|
||||
THREADS_PER_WG = 128
|
||||
NUM_WG = math.ceil(N / TN) * math.ceil(N / TN)
|
||||
|
||||
dname:str = Device.DEFAULT
|
||||
template:str = (pathlib.Path(__file__).parent/"template.s").read_text()
|
||||
|
||||
def asm_kernel() -> UOp:
|
||||
lidx = UOp.special(THREADS_PER_WG, "lidx0")
|
||||
gidx = UOp.special(NUM_WG, "gidx0")
|
||||
|
||||
a = UOp.placeholder((N*N,), dtypes.half, slot=1)
|
||||
b = UOp.placeholder((N*N,), dtypes.half, slot=2)
|
||||
c = UOp.placeholder((N*N,), dtypes.half, slot=0)
|
||||
|
||||
src = template.replace("INSTRUCTIONS", (pathlib.Path(__file__).parent/"gemm.s").read_text())
|
||||
|
||||
sink = UOp.sink(a, b, c, lidx, gidx, arg=KernelInfo(name="gemm"))
|
||||
return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=dname), UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=src)))
|
||||
|
||||
if __name__ == "__main__":
|
||||
test_matmul(asm_kernel(), dtype=dtypes.half, N=N)
|
||||
@@ -1,179 +0,0 @@
|
||||
# unpack the complete kernel descriptor of an amdgpu ELF
|
||||
# https://rocm.docs.amd.com/projects/llvm-project/en/latest/LLVM/llvm/html/AMDGPUUsage.html#code-object-v3-kernel-descriptor
|
||||
import struct, pathlib, sys
|
||||
from tinygrad.runtime.support.elf import elf_loader
|
||||
|
||||
def bits(x, lo, hi): return (x >> lo) & ((1 << (hi - lo + 1)) - 1)
|
||||
def assert_zero(x, lo, hi): assert bits(x, lo, hi) == 0
|
||||
|
||||
with open(sys.argv[1], "rb") as f:
|
||||
lib = f.read()
|
||||
|
||||
image, sections, relocs = elf_loader(lib)
|
||||
rodata_entry = next((sh.header.sh_addr for sh in sections if sh.name == ".rodata"))
|
||||
|
||||
# rodata is exactly 64 bytes
|
||||
kd = image[rodata_entry:rodata_entry+64]
|
||||
desc = int.from_bytes(kd, byteorder="little")
|
||||
|
||||
group_segment_fixed_size = bits(desc, 0, 31)
|
||||
private_segment_fixed_size = bits(desc, 32, 63)
|
||||
kernarg_size = bits(desc, 64, 95)
|
||||
reserved_127_96 = bits(desc, 96, 127)
|
||||
assert reserved_127_96 == 0
|
||||
|
||||
print("GROUP_SEGMENT_FIXED_SIZE:", group_segment_fixed_size)
|
||||
print("PRIVATE_SEGMENT_FIXED_SIZE:", private_segment_fixed_size)
|
||||
print("KERNARG_SIZE:", kernarg_size)
|
||||
print("RESERVED 127:96:", reserved_127_96)
|
||||
|
||||
entry_off = bits(desc, 128, 191)
|
||||
|
||||
# sign-extend manually if needed
|
||||
if entry_off & (1 << 63):
|
||||
entry_off -= 1 << 64
|
||||
|
||||
print("KERNEL_CODE_ENTRY_BYTE_OFFSET:", entry_off)
|
||||
|
||||
kd_addr = 0x1840
|
||||
entry_addr = kd_addr + entry_off
|
||||
|
||||
print("Computed entry address: 0x%016x" % entry_addr)
|
||||
print("256B aligned:", entry_addr % 256 == 0)
|
||||
|
||||
pgm_rsrc3 = bits(desc, 352, 383)
|
||||
pgm_rsrc1 = bits(desc, 384, 415)
|
||||
pgm_rsrc2 = bits(desc, 416, 447)
|
||||
|
||||
print("COMPUTE_PGM_RSRC3: 0x%08x" % pgm_rsrc3)
|
||||
print("COMPUTE_PGM_RSRC1: 0x%08x" % pgm_rsrc1)
|
||||
print("COMPUTE_PGM_RSRC2: 0x%08x" % pgm_rsrc2)
|
||||
|
||||
# rsrc 3 (gfx950)
|
||||
|
||||
accum_offset_raw = bits(pgm_rsrc3, 0, 5)
|
||||
assert_zero(pgm_rsrc3, 6, 15)
|
||||
tg_split = bits(pgm_rsrc3, 16, 16)
|
||||
accum_offset_vgprs = (accum_offset_raw + 1) * 4
|
||||
print("RSRC3.ACCUM_OFFSET (AccVGPR index):", accum_offset_vgprs)
|
||||
print("RSRC3.TG_SPLIT:", tg_split)
|
||||
|
||||
# rsrc 1
|
||||
|
||||
vgpr_gran = bits(pgm_rsrc1, 0, 5)
|
||||
sgpr_gran = bits(pgm_rsrc1, 6, 9)
|
||||
assert_zero(pgm_rsrc1, 27, 28)
|
||||
|
||||
# NOTE: this is vgprs + agprs
|
||||
vgprs_used = (vgpr_gran + 1) * 8
|
||||
assert 0 <= vgprs_used <= 512
|
||||
|
||||
k = sgpr_gran // 2
|
||||
sgprs_used = (k + 1) * 16
|
||||
|
||||
print("RSRC1.VGPRS:", vgprs_used)
|
||||
print("RSRC1.SGPRS:", sgprs_used)
|
||||
|
||||
assert_zero(pgm_rsrc1, 10, 11)
|
||||
|
||||
float_round_mode_32 = bits(pgm_rsrc1, 12, 13)
|
||||
float_round_mode_16_64 = bits(pgm_rsrc1, 15, 14)
|
||||
float_denorm_mode_32 = bits(pgm_rsrc1, 16, 17)
|
||||
float_denorm_mode_16_64 = bits(pgm_rsrc1, 18, 19)
|
||||
|
||||
priv = bits(pgm_rsrc1, 20, 20)
|
||||
assert priv == 0
|
||||
enable_dx10_clamp_wg_rr_en = bits(pgm_rsrc1, 21, 21)
|
||||
debug_mode = bits(pgm_rsrc1, 22, 22)
|
||||
enable_ieee_mode = bits(pgm_rsrc1, 23, 23)
|
||||
bulky = bits(pgm_rsrc1, 24, 24)
|
||||
assert bulky == 0
|
||||
cdbg_user = bits(pgm_rsrc1, 25, 25)
|
||||
assert cdbg_user == 0
|
||||
fp16_ovfl = bits(pgm_rsrc1, 26, 26)
|
||||
assert_zero(pgm_rsrc1, 27, 28) # reserved
|
||||
assert_zero(pgm_rsrc1, 29, 29) # WGP_MODE (reserved on gfx9)
|
||||
assert_zero(pgm_rsrc1, 30, 30) # MEM_ORDERED (reserved on gfx9)
|
||||
assert_zero(pgm_rsrc1, 31, 31) # FWD_PROGRESS (reserved on gfx9)
|
||||
|
||||
# rsrc 2
|
||||
|
||||
enable_private_segment = bits(pgm_rsrc2, 0, 0) # SCRATCH_EN
|
||||
user_sgpr_count = bits(pgm_rsrc2, 1, 5) # USER_SGPR
|
||||
enable_trap_handler = bits(pgm_rsrc2, 6, 6) # TRAP_PRESENT (must be 0 here)
|
||||
assert enable_trap_handler == 0
|
||||
|
||||
enable_sgpr_workgroup_id_x = bits(pgm_rsrc2, 7, 7)
|
||||
enable_sgpr_workgroup_id_y = bits(pgm_rsrc2, 8, 8)
|
||||
enable_sgpr_workgroup_id_z = bits(pgm_rsrc2, 9, 9)
|
||||
enable_sgpr_workgroup_info = bits(pgm_rsrc2, 10, 10)
|
||||
|
||||
enable_vgpr_workitem_id = bits(pgm_rsrc2, 11, 12) # TIDIG_CMP_CNT enum (0..3)
|
||||
|
||||
enable_exception_address_watch = bits(pgm_rsrc2, 13, 13)
|
||||
assert enable_exception_address_watch == 0
|
||||
enable_exception_memory = bits(pgm_rsrc2, 14, 14)
|
||||
assert enable_exception_memory == 0
|
||||
|
||||
granulated_lds_size = bits(pgm_rsrc2, 15, 23)
|
||||
assert granulated_lds_size == 0 # spec: must be 0; CP uses dispatch packet rounding
|
||||
|
||||
enable_exception_fp_invalid = bits(pgm_rsrc2, 24, 24)
|
||||
enable_exception_fp_denorm_src = bits(pgm_rsrc2, 25, 25)
|
||||
enable_exception_fp_div0 = bits(pgm_rsrc2, 26, 26)
|
||||
enable_exception_fp_overflow = bits(pgm_rsrc2, 27, 27)
|
||||
enable_exception_fp_underflow = bits(pgm_rsrc2, 28, 28)
|
||||
enable_exception_fp_inexact = bits(pgm_rsrc2, 29, 29)
|
||||
enable_exception_int_div0 = bits(pgm_rsrc2, 30, 30)
|
||||
|
||||
assert_zero(pgm_rsrc2, 31, 31)
|
||||
|
||||
print("RSRC2.ENABLE_PRIVATE_SEGMENT:", enable_private_segment)
|
||||
print("RSRC2.USER_SGPR_COUNT:", user_sgpr_count)
|
||||
print("RSRC2.ENABLE_SGPR_WORKGROUP_ID_X:", enable_sgpr_workgroup_id_x)
|
||||
print("RSRC2.ENABLE_SGPR_WORKGROUP_ID_Y:", enable_sgpr_workgroup_id_y)
|
||||
print("RSRC2.ENABLE_SGPR_WORKGROUP_ID_Z:", enable_sgpr_workgroup_id_z)
|
||||
print("RSRC2.ENABLE_SGPR_WORKGROUP_INFO:", enable_sgpr_workgroup_info)
|
||||
print("RSRC2.ENABLE_VGPR_WORKITEM_ID (enum):", enable_vgpr_workitem_id)
|
||||
|
||||
print("RSRC2.EXC_FP_INVALID:", enable_exception_fp_invalid)
|
||||
print("RSRC2.EXC_FP_DENORM_SRC:", enable_exception_fp_denorm_src)
|
||||
print("RSRC2.EXC_FP_DIV0:", enable_exception_fp_div0)
|
||||
print("RSRC2.EXC_FP_OVERFLOW:", enable_exception_fp_overflow)
|
||||
print("RSRC2.EXC_FP_UNDERFLOW:", enable_exception_fp_underflow)
|
||||
print("RSRC2.EXC_FP_INEXACT:", enable_exception_fp_inexact)
|
||||
print("RSRC2.EXC_INT_DIV0:", enable_exception_int_div0)
|
||||
|
||||
# user sgprs
|
||||
|
||||
enable_sgpr_private_segment_buffer = bits(desc, 448, 448)
|
||||
enable_sgpr_dispatch_ptr = bits(desc, 449, 449)
|
||||
enable_sgpr_queue_ptr = bits(desc, 450, 450)
|
||||
enable_sgpr_kernarg_segment_ptr = bits(desc, 451, 451)
|
||||
enable_sgpr_dispatch_id = bits(desc, 452, 452)
|
||||
enable_sgpr_flat_scratch_init = bits(desc, 453, 453)
|
||||
enable_sgpr_private_segment_size = bits(desc, 454, 454)
|
||||
|
||||
assert_zero(desc, 455, 457)
|
||||
|
||||
print("DESC.ENABLE_SGPR_PRIVATE_SEGMENT_BUFFER:", enable_sgpr_private_segment_buffer)
|
||||
print("DESC.ENABLE_SGPR_DISPATCH_PTR:", enable_sgpr_dispatch_ptr)
|
||||
print("DESC.ENABLE_SGPR_QUEUE_PTR:", enable_sgpr_queue_ptr)
|
||||
print("DESC.ENABLE_SGPR_KERNARG_SEGMENT_PTR:", enable_sgpr_kernarg_segment_ptr)
|
||||
print("DESC.ENABLE_SGPR_DISPATCH_ID:", enable_sgpr_dispatch_id)
|
||||
print("DESC.ENABLE_SGPR_FLAT_SCRATCH_INIT:", enable_sgpr_flat_scratch_init)
|
||||
print("DESC.ENABLE_SGPR_PRIVATE_SEGMENT_SIZE:", enable_sgpr_private_segment_size)
|
||||
|
||||
assert_zero(desc, 458, 459)
|
||||
|
||||
uses_dynamic_stack = bits(desc, 459, 460)
|
||||
print("DESC.USES_DYNAMIC_STACK:", uses_dynamic_stack)
|
||||
|
||||
# gfx950 only
|
||||
assert_zero(desc, 460, 463)
|
||||
kernarg_preload_spec_length = bits(desc, 464, 470)
|
||||
print("DESC.KERNARG_PRELOAD_SPEC_LENGTH:", kernarg_preload_spec_length)
|
||||
kernarg_preload_spec_offset = bits(desc, 471, 479)
|
||||
print("DESC.KERNARG_PRELOAD_SPEC_OFFSET:", kernarg_preload_spec_offset)
|
||||
|
||||
assert_zero(desc, 480, 511)
|
||||
@@ -1,5 +1,14 @@
|
||||
import atexit, functools
|
||||
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
|
||||
from tinygrad.runtime.autogen.amd.cdna.ins import *
|
||||
from tinygrad.dtype import dtypes
|
||||
|
||||
# ** CDNA4 assembly gemm
|
||||
|
||||
WORKGROUP_SIZE = 256
|
||||
|
||||
# M0 is encoded with 124 (NULL in RDNA) in CDNA
|
||||
M0 = NULL
|
||||
@@ -2597,3 +2606,117 @@ def build_kernel(batch, M, N, K, dtype):
|
||||
k.label('KernelEnd')
|
||||
k.emit(s_endpgm())
|
||||
return k.finalize()
|
||||
|
||||
# ** ASM_GEMM custom kernel
|
||||
|
||||
@functools.cache
|
||||
def custom_asm_gemm(C:UOp, A:UOp, B:UOp, dname:str) -> UOp:
|
||||
batch, M, K = A.shape
|
||||
K2, N = B.shape[(1 if B.ndim == 3 else 0):]
|
||||
assert K == K2
|
||||
lidx = UOp.special(WORKGROUP_SIZE, "lidx0")
|
||||
gidx = UOp.special(NUM_WG, "gidx0")
|
||||
insts = build_kernel(batch, M, N, K, A.dtype.base)
|
||||
lds = UOp(Ops.DEFINE_LOCAL, dtypes.uint8.ptr(size=133_120, addrspace=AddrSpace.LOCAL), (), 'lds')
|
||||
sink = UOp.sink(C.base, A.base, B.base, lds, lidx, gidx,
|
||||
arg=KernelInfo(name=f"gemm_{batch}_{M}_{N}_{K}", estimates=Estimates(ops=2*batch*M*N*K, mem=(batch*M*K + K*N + batch*M*N)*2)))
|
||||
return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=dname),
|
||||
UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in insts]))))
|
||||
|
||||
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}: return todo(f"only bfloat16/float16, 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 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 = getattr(Device[dname].renderer, "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, AxisType.LOOP)
|
||||
n = UOp.range(N, 2, AxisType.LOOP)
|
||||
k = UOp.range(K, 0, AxisType.REDUCE)
|
||||
mul = (A.flatten().index((m*UOp.const(dtypes.index, K)+k))*
|
||||
B.flatten().index((k*UOp.const(dtypes.index, N)+n))).cast(dtypes.float32)
|
||||
red = mul.reduce(k, arg=Ops.ADD, dtype=dtypes.float32).cast(C.dtype.base)
|
||||
store = C.flatten().index((m*UOp.const(dtypes.index, N)+n), ptr=True).store(red).end(m, n)
|
||||
return store.sink(arg=KernelInfo(name=f'uop_gemm_{M}_{N}_{K}'))
|
||||
|
||||
# ** backward gemm, might use the asm gemm
|
||||
|
||||
def custom_gemm_bw(gradient:UOp, kernel:UOp):
|
||||
out, a, b = kernel.src[1:]
|
||||
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)
|
||||
# TODO: this needs to be cleaned up and done properly, the batch dim of grad and a multi need to align
|
||||
g_t = g_t[:a.shape[0]]
|
||||
grad_a = (g_t @ b_t.T).uop
|
||||
grad_b = (a_t.permute(2, 0, 1).reshape(a_t.shape[2], -1) @ g_t.reshape(-1, g_t.shape[-1])).uop
|
||||
return (None, grad_a, grad_b)
|
||||
|
||||
# ** main gemm function
|
||||
|
||||
def asm_gemm(a:Tensor, b:Tensor) -> Tensor:
|
||||
assert can_use_asm_gemm(a, b), f"{counters['todos'][-1]}"
|
||||
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)
|
||||
|
||||
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.empty(batch, M, N//len(a.device), dtype=a.dtype, device=a.device).uop.multi(2), device=a.device)
|
||||
elif m_sharded:
|
||||
out = Tensor(Tensor.empty(batch, M, N, dtype=a.dtype, device=a.device).uop.multi(1), device=a.device)
|
||||
else:
|
||||
out = Tensor(Tensor.empty(batch//len(a.device) if a.uop.axis==0 else batch, M, N, dtype=a.dtype, device=a.device).uop.multi(0), device=a.device)
|
||||
else:
|
||||
out = Tensor.empty(batch, M, N, dtype=a.dtype, device=a.device)
|
||||
|
||||
renderer = Device[a.device[0] if is_multi else a.device].renderer
|
||||
dname, arch = renderer.device, getattr(renderer, "arch", "")
|
||||
if arch.startswith("gfx950") and getenv("USE_ASM", 1):
|
||||
out = Tensor.custom_kernel(out, a, b, fxn=functools.partial(custom_asm_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])
|
||||
return out
|
||||
@@ -2,7 +2,7 @@ import unittest
|
||||
from tinygrad import Tensor, Device, dtypes, Context
|
||||
from tinygrad.device import is_dtype_supported
|
||||
from tinygrad.helpers import getenv
|
||||
from extra.gemm.asm.cdna.gemm import asm_gemm
|
||||
from extra.gemm.cdna_asm_gemm import asm_gemm
|
||||
from test.helpers import needs_second_gpu
|
||||
|
||||
# On non CDNA4 it will only validate the Tensor.custom_kernel integration
|
||||
@@ -157,7 +157,7 @@ class TestGemmLarge(unittest.TestCase):
|
||||
|
||||
class TestMagicGu(unittest.TestCase):
|
||||
def test_magicgu_matches_old(self):
|
||||
from extra.gemm.asm.cdna.asm import _magicgu_mulhi, TILE_M, TILE_N, TILE_K
|
||||
from extra.gemm.cdna_asm_gemm import _magicgu_mulhi, TILE_M, TILE_N, TILE_K
|
||||
old_iters_args = {64: (67108864, 0), 128: (33554432, 0), 224: (613566757, 2147483656)}
|
||||
old_gemm_shapes = [
|
||||
(8192, 4096, 4096), (8192, 14336, 4096), (8192, 4096, 14336),
|
||||
|
||||
+1
-1
@@ -2491,7 +2491,7 @@ class Tensor(OpMixin):
|
||||
"""
|
||||
if IMAGE: return self.image_dot(w, dtype)
|
||||
if ASM_GEMM:
|
||||
from extra.gemm.asm.cdna.gemm import can_use_asm_gemm, asm_gemm
|
||||
from extra.gemm.cdna_asm_gemm import can_use_asm_gemm, asm_gemm
|
||||
if can_use_asm_gemm(self, w): return asm_gemm(self, w)
|
||||
x, dx, dw = self, self.ndim, w.ndim
|
||||
if not (dx > 0 and dw > 0): raise RuntimeError(f"both tensors need to be at least 1D, got {dx}D and {dw}D")
|
||||
|
||||
Reference in New Issue
Block a user