mirror of
https://github.com/tinygrad/tinygrad.git
synced 2026-08-29 13:16:09 +00:00
mmapeak implementation for 7900 XTX (#10417)
* Add mmapeak implementation for 7900 XTX * Change identation * Use a template instead of multiple assebly files * Fix output formatting * Reduce register file bank conflicts * More accurate measurement for quick instructions * Add support for gfx1201 * RDNA4 wmma requires less VGRPs * RDNA4 does not have s_cmpk instructions * Add v_wmma_i32_16x16x32_iu4 for gfx1201 * Add sparse wmma instructions --------- Co-authored-by: George Hotz <[email protected]>
This commit is contained in:
Panagiotis Kourouklidis
co-authored by
George Hotz
George Hotz
parent
0a313d98a0
commit
e21836952d
@@ -0,0 +1,79 @@
|
||||
import pathlib
|
||||
from tinygrad.device import Device
|
||||
from tinygrad.runtime.ops_amd import AMDProgram, HIPCompiler
|
||||
import time
|
||||
import os
|
||||
|
||||
NUM_WORKGROUPS = 96
|
||||
WAVE_SIZE = 32
|
||||
NUM_WAVES = 2
|
||||
FLOPS_PER_MATMUL = 16*16*16*2
|
||||
INTERNAL_LOOP = 1_000_000
|
||||
INSTRUCTIONS_PER_LOOP = 1_000
|
||||
|
||||
assemblyTemplate = (pathlib.Path(__file__).parent / "template.s").read_text()
|
||||
|
||||
def launchBenchmark(instruction, vgprIndices, dense = True):
|
||||
if dense:
|
||||
instructions = "{} v[0:{}], v[{}:{}], v[{}:{}], 1\n".format(instruction, vgprIndices[0],
|
||||
vgprIndices[1], vgprIndices[2],
|
||||
vgprIndices[1], vgprIndices[2]) * INSTRUCTIONS_PER_LOOP
|
||||
else:
|
||||
instructions = "{} v[0:{}], v[{}:{}], v[{}:{}], v{}\n".format(instruction, vgprIndices[0],
|
||||
vgprIndices[1], vgprIndices[2],
|
||||
vgprIndices[3], vgprIndices[4],
|
||||
vgprIndices[5]) * INSTRUCTIONS_PER_LOOP
|
||||
src = assemblyTemplate.replace("INSTRUCTION", instructions)
|
||||
lib = COMPILER.compile(src)
|
||||
fxn = AMDProgram(DEV, "matmul", lib)
|
||||
start = time.perf_counter()
|
||||
fxn(global_size=(NUM_WORKGROUPS,1,1), local_size=(WAVE_SIZE*NUM_WAVES,1,1), wait=True) #For some reason the returned time is very small after the first kernel execution
|
||||
end = time.perf_counter()
|
||||
elapsed = end-start
|
||||
FLOPs = FLOPS_PER_MATMUL * NUM_WAVES * NUM_WORKGROUPS * INTERNAL_LOOP * INSTRUCTIONS_PER_LOOP
|
||||
print("{:<29} : {} T(FL)OPS".format(instruction, round(FLOPs/elapsed/10**12, 2)))
|
||||
|
||||
if __name__=="__main__":
|
||||
DEVICENUM = os.getenv("DEVICENUM", "0")
|
||||
try:
|
||||
DEV = Device['AMD:' + DEVICENUM]
|
||||
except:
|
||||
raise RuntimeError("Error while initiating AMD device")
|
||||
|
||||
if (ARCH := DEV.arch) not in ['gfx1100', 'gfx1201']:
|
||||
raise RuntimeError("only gfx1100 and gfx1201 supported")
|
||||
COMPILER = HIPCompiler(ARCH)
|
||||
|
||||
if ARCH == 'gfx1100':
|
||||
launchBenchmark("v_wmma_bf16_16x16x16_bf16", (7,8,15))
|
||||
launchBenchmark("v_wmma_f16_16x16x16_f16", (7,8,15))
|
||||
launchBenchmark("v_wmma_f32_16x16x16_bf16", (7,8,15))
|
||||
launchBenchmark("v_wmma_f32_16x16x16_f16", (7,8,15))
|
||||
launchBenchmark("v_wmma_i32_16x16x16_iu4", (7,8,9))
|
||||
launchBenchmark("v_wmma_i32_16x16x16_iu8", (7,8,11))
|
||||
if ARCH == 'gfx1201':
|
||||
NUM_WORKGROUPS = 64
|
||||
launchBenchmark("v_wmma_bf16_16x16x16_bf16", (3,4,7))
|
||||
launchBenchmark("v_wmma_f16_16x16x16_f16", (3,4,7))
|
||||
launchBenchmark("v_wmma_f32_16x16x16_bf16", (7,8,11))
|
||||
launchBenchmark("v_wmma_f32_16x16x16_f16", (7,8,11))
|
||||
launchBenchmark("v_wmma_i32_16x16x16_iu4", (7,8,8))
|
||||
launchBenchmark("v_wmma_i32_16x16x16_iu8", (7,8,9))
|
||||
launchBenchmark("v_wmma_f32_16x16x16_fp8_fp8", (7,8,9))
|
||||
launchBenchmark("v_wmma_f32_16x16x16_fp8_bf8", (7,8,9))
|
||||
launchBenchmark("v_wmma_f32_16x16x16_bf8_fp8", (7,8,9))
|
||||
launchBenchmark("v_wmma_f32_16x16x16_bf8_bf8", (7,8,9))
|
||||
FLOPS_PER_MATMUL = 16*16*32*2
|
||||
launchBenchmark("v_wmma_i32_16X16X32_iu4", (7,8,9))
|
||||
launchBenchmark("v_swmmac_f32_16x16x32_f16", (7,8,11,12,19,20), False)
|
||||
launchBenchmark("v_swmmac_f32_16x16x32_bf16", (7,8,11,12,19,20), False)
|
||||
launchBenchmark("v_swmmac_f16_16x16x32_f16", (3,4,7,8,15,16), False)
|
||||
launchBenchmark("v_swmmac_bf16_16x16x32_bf16", (3,4,7,8,15,16), False)
|
||||
launchBenchmark("v_swmmac_i32_16x16x32_iu8", (7,8,9,10,13,14), False)
|
||||
launchBenchmark("v_swmmac_i32_16x16x32_iu4", (7,8,8,9,10,11), False)
|
||||
launchBenchmark("v_swmmac_f32_16x16x32_fp8_fp8", (7,8,9,10,13,14), False)
|
||||
launchBenchmark("v_swmmac_f32_16x16x32_fp8_bf8", (7,8,9,10,13,14), False)
|
||||
launchBenchmark("v_swmmac_f32_16x16x32_bf8_fp8", (7,8,9,10,13,14), False)
|
||||
launchBenchmark("v_swmmac_f32_16x16x32_bf8_bf8", (7,8,9,10,13,14), False)
|
||||
FLOPS_PER_MATMUL = 16*16*64*2
|
||||
launchBenchmark("v_swmmac_i32_16x16x64_iu4", (7,8,9,10,13,14), False)
|
||||
@@ -0,0 +1,41 @@
|
||||
.text
|
||||
.globl matmul
|
||||
.p2align 8
|
||||
.type matmul,@function
|
||||
matmul:
|
||||
s_mov_b32 s1, 1000000
|
||||
s_mov_b32 s2, 0
|
||||
inner_loop:
|
||||
INSTRUCTION
|
||||
s_sub_u32 s1, s1, 1
|
||||
s_cmp_lg_i32 s1, s2
|
||||
s_cbranch_scc1 inner_loop
|
||||
s_endpgm
|
||||
|
||||
.rodata
|
||||
.p2align 6
|
||||
.amdhsa_kernel matmul
|
||||
.amdhsa_next_free_vgpr .amdgcn.next_free_vgpr
|
||||
.amdhsa_next_free_sgpr .amdgcn.next_free_sgpr
|
||||
.amdhsa_wavefront_size32 1
|
||||
.end_amdhsa_kernel
|
||||
|
||||
.amdgpu_metadata
|
||||
---
|
||||
amdhsa.version:
|
||||
- 1
|
||||
- 0
|
||||
amdhsa.kernels:
|
||||
- .name: matmul
|
||||
.symbol: matmul.kd
|
||||
.kernarg_segment_size: 0
|
||||
.group_segment_fixed_size: 0
|
||||
.private_segment_fixed_size: 0
|
||||
.kernarg_segment_align: 4
|
||||
.wavefront_size: 32
|
||||
.sgpr_count: 8
|
||||
.vgpr_count: 32
|
||||
.max_flat_workgroup_size: 1024
|
||||
.args:
|
||||
...
|
||||
.end_amdgpu_metadata
|
||||
Reference in New Issue
Block a user