derek.holloway
7502bf64a4
CodeQL / Analyze (actions) (push) Canceled after 0s
CodeQL / Analyze (python) (push) Canceled after 0s
pre-commit / pre-commit (push) Canceled after 0s
Multi-Arch CI / setup (push) Canceled after 0s
Publish build_manylinux_x86_64 images / publish_build_manylinux_x86_64 (push) Canceled after 0s
Setup (Multi-Arch) / setup (push) Canceled after 0s
Multi-Arch CI / Linux::${{ fromJSON(needs.setup.outputs.linux_build_config || '{}').build_variant_label || 'skip' }} (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Publish no_rocm_image_ubi10 image / publish_no_rocm_image_ubi10 (push) Canceled after 0s
Multi-Arch CI / Windows::${{ fromJSON(needs.setup.outputs.windows_build_config || '{}').build_variant_label || 'skip' }} (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Multi-Arch CI / Manifest Diff (push) Canceled after 0s
TheRock Manifest Diff Report / Generate Manifest Diff Report (push) Canceled after 0s
Multi-Arch CI / CI Summary (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04 (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04_ocl_rt images / publish_no_rocm_image_ubuntu24_04_ocl_rt (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_rocgdb (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04_rocgdb images / publish_no_rocm_image_ubuntu24_04_rocgdb (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_ocl_rt (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Unit Tests / Unit Tests :: ubuntu-24.04 (push) Canceled after 0s
Unit Tests / Unit Tests :: windows-2022 (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_media (push) Canceled after 0s
Unit Tests / Unit Tests Summary (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_openmpi (push) Canceled after 0s
909 lines
39 KiB
Python
909 lines
39 KiB
Python
################################################################################
|
|
#
|
|
# Copyright (C) 2023-2025 Advanced Micro Devices, Inc. All rights reserved.
|
|
#
|
|
# Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
# of this software and associated documentation files (the "Software"), to deal
|
|
# in the Software without restriction, including without limitation the rights
|
|
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell cop-
|
|
# ies of the Software, and to permit persons to whom the Software is furnished
|
|
# to do so, subject to the following conditions:
|
|
#
|
|
# The above copyright notice and this permission notice shall be included in all
|
|
# copies or substantial portions of the Software.
|
|
#
|
|
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IM-
|
|
# PLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS
|
|
# FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
|
|
# COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER
|
|
# IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNE-
|
|
# CTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
|
|
################################################################################
|
|
|
|
from rocisa.code import Label, Module, RegSet, TextBlock, ValueSet, SrdUpperValue
|
|
from rocisa.container import EXEC, VCC, DSModifiers, MUBUFModifiers, vgpr, sgpr
|
|
from rocisa.enum import RegisterType
|
|
from rocisa.register import RegisterPool
|
|
import rocisa.instruction as ri
|
|
|
|
from argparse import ArgumentParser
|
|
from dataclasses import dataclass
|
|
from typing import List, Tuple, Optional
|
|
from math import log2
|
|
import os
|
|
import yaml
|
|
import json
|
|
import collections
|
|
from contextlib import contextmanager
|
|
from Tensile.Common.Utilities import _global_ti
|
|
from Tensile.Common.Architectures import detectGlobalCurrentISA, isaToGfx, gfxToIsa
|
|
from Tensile.Common.DataType import DataType
|
|
from Tensile.Common.GlobalParameters import restoreDefaultGlobalParameters, assignGlobalParameters
|
|
from Tensile.Common.Types import IsaVersion
|
|
from Tensile.Toolchain.Validators import ToolchainDefaults, validateToolchain
|
|
|
|
def kernel_header(name: str, gfx_arch: str, vgpr: int, sgpr: int, lds: int, xnack: bool = False):
|
|
vgpr = ((vgpr+7)//8)*8
|
|
sgpr = ((sgpr+7)//8)*8
|
|
lds = ((lds+31)//32)*32
|
|
|
|
target_id = f'{gfx_arch}:xnack+' if xnack else gfx_arch
|
|
header = ""
|
|
header += f'.amdgcn_target "amdgcn-amd-amdhsa--{target_id}"\n'
|
|
header += f'.text\n'
|
|
header += f'.protected {name}\n'
|
|
header += f'.globl {name}\n'
|
|
header += f'.p2align 8\n'
|
|
header += f'.type {name},@function\n'
|
|
header += f'.section .rodata,#alloc\n'
|
|
header += f'.p2align 6\n'
|
|
header += f'.amdhsa_kernel {name}\n'
|
|
header += f' .amdhsa_user_sgpr_kernarg_segment_ptr 1\n'
|
|
if (gfx_arch not in ("gfx900", "gfx908", "gfx1030", "gfx1100", "gfx1101", "gfx1102", "gfx1103", "gfx1150", "gfx1151", "gfx1152", "gfx1153", "gfx1200", "gfx1201", "gfx1250")):
|
|
header += f' .amdhsa_accum_offset {vgpr} // accvgpr offset\n'
|
|
header += f' .amdhsa_next_free_vgpr {vgpr} // vgprs\n'
|
|
header += f' .amdhsa_next_free_sgpr {sgpr} // sgprs\n'
|
|
header += f' .amdhsa_group_segment_fixed_size {lds} // lds bytes\n'
|
|
header += f' .amdhsa_private_segment_fixed_size 0\n'
|
|
header += f' .amdhsa_system_sgpr_workgroup_id_x 1\n'
|
|
header += f' .amdhsa_system_sgpr_workgroup_id_y 1\n'
|
|
header += f' .amdhsa_system_sgpr_workgroup_id_z 1\n'
|
|
header += f' .amdhsa_system_vgpr_workitem_id 0\n'
|
|
header += f' .amdhsa_float_denorm_mode_32 3\n'
|
|
header += f' .amdhsa_float_denorm_mode_16_64 3\n'
|
|
if _global_ti.getArchCaps()["HasWave32"]:
|
|
header += f' .amdhsa_wavefront_size32 1\n'
|
|
header += f'.end_amdhsa_kernel\n'
|
|
header += f'.text\n'
|
|
return header
|
|
|
|
@contextmanager
|
|
def asm_func(func_name: str, module: Module):
|
|
try:
|
|
module.add(TextBlock(f'{func_name}:\n'))
|
|
yield
|
|
finally:
|
|
end_label_name = f'.L{func_name}_end'
|
|
module.add(ri.SEndpgm())
|
|
module.add(TextBlock(f'{end_label_name}:\n'))
|
|
module.add(TextBlock(f'.size {func_name}, {end_label_name} - {func_name}\n'))
|
|
|
|
@contextmanager
|
|
def asm_loop(mod: Module, name: str, it: str):
|
|
try:
|
|
loop_start_label = Label(name, f'loop {name} starts')
|
|
loop_end_label = Label(f'{name}_end', f'loop {name} ends')
|
|
mod.add(loop_start_label)
|
|
mod.add(ri.SCmpEQU32(sgpr(it), 0))
|
|
mod.add(ri.SCBranchSCC1(loop_end_label.getLabelName()))
|
|
mod.addSpaceLine()
|
|
yield
|
|
finally:
|
|
mod.add(ri.SSubU32(sgpr(it), sgpr(it), 1))
|
|
mod.add(ri.SBranch(loop_start_label.getLabelName()))
|
|
mod.add(loop_end_label)
|
|
mod.addSpaceLine()
|
|
|
|
|
|
class AMaxKernelGenerator:
|
|
srd_num_reg = 4
|
|
srd_alignment = 4
|
|
|
|
def __init__(self,
|
|
i_type: DataType,
|
|
o_type: DataType,
|
|
scale_type: DataType,
|
|
num_workitems: int,
|
|
num_load_count: int,
|
|
num_load_size: int,
|
|
wavefront_size: int,
|
|
arch: str,
|
|
isa: IsaVersion,
|
|
is_scale: bool):
|
|
self.i_type = i_type
|
|
self.o_type = o_type
|
|
self.scale_type = scale_type
|
|
self.bpe = int(i_type.numBytes())
|
|
self.num_workitems = num_workitems
|
|
self.num_load_count = num_load_count
|
|
self.num_load_size = num_load_size
|
|
self.wavefront_size = wavefront_size
|
|
self.laneSGPRCount = 2 if wavefront_size == 64 else 1
|
|
self.SMovBX = ri.SMovB64 if wavefront_size == 64 else ri.SMovB32
|
|
self.SAndBX = ri.SAndB64 if wavefront_size == 64 else ri.SAndB32
|
|
self.srdUpperValue = SrdUpperValue(isa)
|
|
self.sgpr_pool = RegisterPool(24, RegisterType.Sgpr, True)
|
|
self.vgpr_pool = RegisterPool(40, RegisterType.Vgpr, True)
|
|
self.sgpr_pool.add(0, 23) #TODO: estimate this
|
|
self.vgpr_pool.add(0, 39) #TODO: estimate this
|
|
self.debug_label = True
|
|
self.arch = arch
|
|
self.isa = isa
|
|
self.is_scale = is_scale
|
|
self.op = 'AMax'
|
|
self.sgprs = collections.OrderedDict()
|
|
self.vgprs = collections.OrderedDict()
|
|
|
|
@property
|
|
def lds_usage_byte(self) -> int:
|
|
# used in reduce inter wave mean and invvar
|
|
# 4 data * half_wave_num * bpe
|
|
return 4 * (self.num_workitems // self.wavefront_size // 2) * self.bpe
|
|
|
|
@property
|
|
def func_name(self):
|
|
if self.is_scale:
|
|
return f'AMax_Ti_{self.i_type}_To_{self.o_type}_Ts_{self.scale_type}_W_{self.num_workitems}_C_{self.num_load_count}'
|
|
return f'AMax_Ti_{self.i_type}_To_{self.o_type}_W_{self.num_workitems}_C_{self.num_load_count}'
|
|
|
|
def dumps(self, format: str) -> str:
|
|
param_dict = {
|
|
'arch': self.arch,
|
|
'op': self.op,
|
|
'is_scale': self.is_scale,
|
|
'func_name': self.func_name,
|
|
'io_type': self.i_type.toChar(),
|
|
'o_type': self.o_type.toChar(),
|
|
'scale_type': self.scale_type.toChar(),
|
|
'num_workitems': self.num_workitems,
|
|
}
|
|
|
|
if format.lower() == 'yaml':
|
|
return yaml.dump(param_dict)
|
|
elif format.lower() == 'json':
|
|
return json.dumps(param_dict)
|
|
else:
|
|
assert False, f'Unsupported format {format}'
|
|
|
|
def dump(self, format: str, output_path: str):
|
|
s = self.dumps(format)
|
|
with open(output_path, 'w') as f:
|
|
f.write(s)
|
|
|
|
|
|
def local_read_inst_type(self, num_elements: int):
|
|
if self.i_type.isSingle():
|
|
insts = {
|
|
1: ri.DSLoadB32,
|
|
2: ri.DSLoadB64,
|
|
4: ri.DSLoadB128
|
|
}
|
|
elif self.i_type.isHalf():
|
|
insts = {
|
|
1: ri.DSLoadU16,
|
|
2: ri.DSLoadB32,
|
|
4: ri.DSLoadB64,
|
|
8: ri.DSLoadB128
|
|
}
|
|
return insts[num_elements]
|
|
|
|
|
|
def local_write_inst_type(self, num_elements: int):
|
|
if self.i_type.isSingle():
|
|
insts = {
|
|
1: ri.DSStoreB32,
|
|
2: ri.DSStoreB64,
|
|
4: ri.DSStoreB128
|
|
}
|
|
elif self.i_type.isHalf():
|
|
insts = {
|
|
1: ri.DSStoreB16,
|
|
2: ri.DSStoreB32,
|
|
4: ri.DSStoreB64,
|
|
8: ri.DSStoreB128
|
|
}
|
|
return insts[num_elements]
|
|
|
|
|
|
def global_read_inst_type(self, num_elements: int):
|
|
if self.i_type.isSingle():
|
|
insts = {
|
|
1: ri.BufferLoadB32,
|
|
2: ri.BufferLoadB64,
|
|
4: ri.BufferLoadB128
|
|
}
|
|
elif self.i_type.isHalf():
|
|
insts = {
|
|
1: ri.BufferLoadD16B16,
|
|
2: ri.BufferLoadB32,
|
|
4: ri.BufferLoadB64,
|
|
8: ri.BufferLoadB128
|
|
}
|
|
else:
|
|
raise NotImplementedError
|
|
return insts[num_elements]
|
|
|
|
def global_write_inst_type(self, num_elements: int):
|
|
if self.o_type.isSingle():
|
|
insts = {
|
|
1: ri.BufferStoreB32,
|
|
2: ri.BufferStoreB64,
|
|
4: ri.BufferStoreB128
|
|
}
|
|
elif self.o_type.isHalf():
|
|
insts = {
|
|
1: ri.BufferStoreB16,
|
|
2: ri.BufferStoreB32,
|
|
4: ri.BufferStoreB64,
|
|
8: ri.BufferStoreB128
|
|
}
|
|
else:
|
|
raise NotImplementedError
|
|
return insts[num_elements]
|
|
|
|
|
|
|
|
def defineSgpr(self, name, numSgprs, align=1):
|
|
if numSgprs == 0: return
|
|
sgprIdx = self.sgpr_pool.checkOutAligned(numSgprs, align, tag=name, preventOverflow=False)
|
|
self.sgprs[name] = sgprIdx
|
|
return sgprIdx
|
|
|
|
def defineVgpr(self, name, numVgprs, align=1):
|
|
if numVgprs == 0: return
|
|
vgprIdx = self.vgpr_pool.checkOutAligned(numVgprs, align, tag=name, preventOverflow=False)
|
|
self.vgprs[name] = vgprIdx
|
|
return vgprIdx
|
|
|
|
|
|
def kernel_args(self):
|
|
if self.is_scale:
|
|
return (KernelArgument(8, 0, 'global_buffer', 'global'),
|
|
KernelArgument(8, 8, 'global_buffer', 'global'),
|
|
KernelArgument(8, 16, 'global_buffer', 'global'),
|
|
KernelArgument(8, 24, 'global_buffer', 'global'),
|
|
KernelArgument(4, 32, 'by_value'))
|
|
return (KernelArgument(8, 0, 'global_buffer', 'global'),
|
|
KernelArgument(8, 8, 'global_buffer', 'global'),
|
|
KernelArgument(4, 16, 'by_value'))
|
|
|
|
|
|
def defineVariables(self):
|
|
self.defineVgpr("Serial", 1, 1)
|
|
self.defineVgpr("Output", 1, 1)
|
|
self.defineVgpr("OutputB", 1, 1)
|
|
self.defineVgpr("Widx", 1, 1)
|
|
self.defineVgpr("Offset", 4, 1)
|
|
self.defineVgpr("Value", self.num_load_count * self.num_load_size, self.num_load_size)
|
|
self.defineVgpr("Tmp", 4, 1)
|
|
if self.is_scale:
|
|
self.defineVgpr("OffsetD", self.num_load_count * self.num_load_size, 1)
|
|
self.defineVgpr("OutputD", self.num_load_count * self.num_load_size, self.num_load_size)
|
|
self.defineVgpr("TmpD", 4, 1)
|
|
if self.scale_type == DataType("F8") or self.scale_type == DataType("F8N"):
|
|
self.defineVgpr("Fp8NanInf", 1, 1)
|
|
self.defineVgpr("Fp8Max", 1, 1)
|
|
self.defineVgpr("Fp8Min", 1, 1)
|
|
self.defineVgpr("Fp8Tmp", 1, 1)
|
|
elif self.scale_type == DataType("B8") or self.scale_type == DataType("B8N"):
|
|
self.defineVgpr("BF8NanInf", 1, 1)
|
|
self.defineVgpr("BF8Max", 1, 1)
|
|
self.defineVgpr("BF8Min", 1, 1)
|
|
self.defineVgpr("BF8Tmp", 1, 1)
|
|
|
|
self.defineSgpr("KernelArg", 2)
|
|
self.defineSgpr("WorkGroup0", 1)
|
|
self.defineSgpr("WorkGroup1", 1)
|
|
self.defineSgpr("WorkGroup2", 1)
|
|
self.defineSgpr("AddressOut", 2, 2)
|
|
self.defineSgpr("AddressOutD", 2, 2)
|
|
self.defineSgpr("AddressIn", 2, 2)
|
|
self.defineSgpr("AddressScale", 2, 2)
|
|
self.defineSgpr("SizeLength", 1)
|
|
self.defineSgpr("MainLoop", 1)
|
|
self.defineSgpr("Offset", 1)
|
|
self.defineSgpr("Src", 4, 4)
|
|
self.defineSgpr("Dst", 4, 4)
|
|
self.defineSgpr("Tmp", 6, 2)
|
|
if self.is_scale:
|
|
self.defineSgpr("DstD", 4, 4)
|
|
self.defineSgpr("TmpD", 6, 2)
|
|
self.defineSgpr("Scale", 1)
|
|
|
|
mod = Module("defineVariables")
|
|
|
|
for vkey in self.vgprs:
|
|
mod.add(RegSet("v", "vgpr"+vkey, self.vgprs[vkey]))
|
|
mod.addSpaceLine()
|
|
|
|
for skey in self.sgprs:
|
|
mod.add(RegSet("s", "sgpr"+skey, self.sgprs[skey]))
|
|
mod.addSpaceLine()
|
|
|
|
mod.add(ValueSet("Srd127_96", self.srdUpperValue.getValue(), format=1))
|
|
mod.addSpaceLine()
|
|
mod.addSpaceLine()
|
|
return mod
|
|
|
|
|
|
def load_kernel_args(self):
|
|
mod = Module('Load kernel args')
|
|
mod.addComment0('Load kernel args')
|
|
if self.is_scale:
|
|
mod.add(ri.SLoadB64(sgpr("AddressOut", 2), sgpr("KernelArg", 2), 0))
|
|
mod.add(ri.SLoadB64(sgpr("AddressOutD", 2), sgpr("KernelArg", 2), 8))
|
|
mod.add(ri.SLoadB64(sgpr("AddressIn", 2), sgpr("KernelArg", 2), 16))
|
|
mod.add(ri.SLoadB64(sgpr("AddressScale", 2), sgpr("KernelArg", 2), 24))
|
|
mod.add(ri.SLoadB32(sgpr("SizeLength"), sgpr("KernelArg", 2), 32))
|
|
else:
|
|
mod.add(ri.SLoadB64(sgpr("AddressOut", 2), sgpr("KernelArg", 2), 0))
|
|
mod.add(ri.SLoadB64(sgpr("AddressIn", 2), sgpr("KernelArg", 2), 8))
|
|
mod.add(ri.SLoadB32(sgpr("SizeLength"), sgpr("KernelArg", 2), 16))
|
|
mod.add(ri.SWaitCnt(kmcnt=0))
|
|
mod.addSpaceLine()
|
|
mod.addSpaceLine()
|
|
return mod
|
|
|
|
def shiftSrd(self, srdStr):
|
|
module = Module()
|
|
if self.isa[0] == 12 and self.isa[1] == 5:
|
|
stmp = self.sgpr_pool.checkOutAligned(1, 1)
|
|
module.add(ri.SAndB32(sgpr(stmp), sgpr(srdStr+"+2"), 0x7F))
|
|
module.add(ri.SLShiftLeftB32(sgpr(stmp), 25, sgpr(stmp)))
|
|
module.add(ri.SAndB32(sgpr(srdStr+"+1"), sgpr(srdStr+"+1"), 0x1FFFFFF))
|
|
module.add(ri.SOrB32(sgpr(srdStr+"+1"), sgpr(srdStr+"+1"), sgpr(stmp)))
|
|
module.add(ri.SLShiftRightB32(sgpr(srdStr+"+2"), 7, sgpr(srdStr+"+2")))
|
|
self.sgpr_pool.checkIn(stmp)
|
|
return module
|
|
|
|
def init_param(self) -> Module:
|
|
mod = Module("init_param")
|
|
mod.addComment0("init_param")
|
|
mod.add(ri.SLShiftLeftB32(sgpr("Tmp"), int(log2(self.bpe)), sgpr("SizeLength")))
|
|
mod.addSpaceLine()
|
|
|
|
mod.add(ri.SMovB32(sgpr("Dst+0"), sgpr("AddressOut+0")))
|
|
mod.add(ri.SMovB32(sgpr("Dst+1"), sgpr("AddressOut+1")))
|
|
mod.add(ri.SMovB32(sgpr("Dst+2"), int(self.o_type.numBytes())))
|
|
mod.add(ri.SMovB32(sgpr("Dst+3"), "Srd127_96"))
|
|
mod.add(self.shiftSrd("Dst"))
|
|
mod.addSpaceLine()
|
|
|
|
if self.is_scale: # init inputScale
|
|
mod.add(ri.SLoadB32(sgpr("Scale"), sgpr("AddressScale", 2), 0))
|
|
mod.add(ri.SWaitCnt(kmcnt=0))
|
|
mod.addSpaceLine()
|
|
|
|
mod.add(ri.SMovB32(sgpr("Src+0"), sgpr("AddressIn+0")))
|
|
mod.add(ri.SMovB32(sgpr("Src+1"), sgpr("AddressIn+1")))
|
|
mod.add(ri.SMovB32(sgpr("Src+2"), sgpr("Tmp")))
|
|
mod.add(ri.SMovB32(sgpr("Src+3"), "Srd127_96"))
|
|
mod.add(self.shiftSrd("Src"))
|
|
mod.addSpaceLine()
|
|
|
|
mod.add(ri.VMovB32(vgpr("Output"), 0))
|
|
if self.is_scale:
|
|
if self.scale_type == DataType("F8N"):
|
|
mod.add(ri.VMovB32(vgpr("Fp8NanInf"), "0x207", comment="Nan and +/- inf"))
|
|
mod.add(ri.VMovB32(vgpr("Fp8Max"), "0x43700000", comment="Fp8 NANOO Max value 240 as float32"))
|
|
mod.add(ri.VMovB32(vgpr("Fp8Min"), "0xc3700000", comment="Fp8 NANOO Min value -240 as float32"))
|
|
elif self.scale_type == DataType("F8"):
|
|
mod.add(ri.VMovB32(vgpr("Fp8NanInf"), "0x207", comment="Nan and +/- inf"))
|
|
mod.add(ri.VMovB32(vgpr("Fp8Max"), "0x43e00000", comment="Fp8 Max value 448 as float32"))
|
|
mod.add(ri.VMovB32(vgpr("Fp8Min"), "0xc3e00000", comment="Fp8 Min value -448 as float32"))
|
|
elif self.scale_type == DataType("B8") or self.scale_type == DataType("B8N"):
|
|
mod.add(ri.VMovB32(vgpr("BF8NanInf"), "0x207", comment="Nan and +/- inf"))
|
|
mod.add(ri.VMovB32(vgpr("BF8Max"), "0x47600000", comment="BF8 Max value 57344 as float32"))
|
|
mod.add(ri.VMovB32(vgpr("BF8Min"), "0xc7600000", comment="BF8 Min value -57344 as float32"))
|
|
mod.addSpaceLine()
|
|
|
|
if self.is_scale:
|
|
mod.add(ri.SMovB32(sgpr("DstD+0"), sgpr("AddressOutD+0")))
|
|
mod.add(ri.SMovB32(sgpr("DstD+1"), sgpr("AddressOutD+1")))
|
|
mod.add(ri.SMovB32(sgpr("DstD+2"), sgpr("SizeLength")))
|
|
mod.add(ri.SMovB32(sgpr("DstD+3"), "Srd127_96"))
|
|
mod.add(self.shiftSrd("DstD"))
|
|
for i in range(self.num_load_count * self.num_load_size):
|
|
mod.add(ri.VMovB32(vgpr(f"OutputD+{i}"), 0))
|
|
|
|
mod.addSpaceLine()
|
|
return mod
|
|
|
|
|
|
def calculate_global_address(self) -> Module:
|
|
mod = Module("calculate_global_address")
|
|
mod.addComment0("calculate_global_address")
|
|
# offset for buffer load
|
|
# total load size = dwordx4 = 16 bytes per PE
|
|
mod.add(ri.VLShiftLeftB32(vgpr("Offset+0"), hex(int(log2(self.num_load_size * 4))), vgpr("Serial")))
|
|
mod.add(ri.SMovB32(sgpr("Tmp"), self.num_workitems * self.num_load_size * 4))
|
|
for i in range(0, self.num_load_count-1):
|
|
mod.add(ri.VAddU32(vgpr(f"Offset+{i+1}"), vgpr(f"Offset+{i}"), sgpr("Tmp")))
|
|
mod.addSpaceLine()
|
|
|
|
if self.is_scale: # offset for buffer store
|
|
# total store size = 1byte x 4 = 4 bytes per PE
|
|
mod.add(ri.VLShiftLeftB32(vgpr("OffsetD+0"), hex(int(log2(4))), vgpr("Serial")))
|
|
mod.add(ri.SMovB32(sgpr("TmpD"), 1))
|
|
for i in range(self.num_load_size - 1):
|
|
mod.add(ri.VAddU32(vgpr(f"OffsetD+{i+1}"), vgpr(f"OffsetD+{i}"), sgpr("TmpD")))
|
|
mod.add(ri.SMovB32(sgpr("TmpD"), self.num_workitems * 4))
|
|
for i in range(0, self.num_load_count-1):
|
|
for j in range(0, self.num_load_size):
|
|
mod.add(ri.VAddU32(vgpr(f"OffsetD+{j+(i+1)*self.num_load_size}"), \
|
|
vgpr(f"OffsetD+{j+i*self.num_load_size}"), sgpr("TmpD")))
|
|
|
|
mod.addSpaceLine()
|
|
mod.addSpaceLine()
|
|
return mod
|
|
|
|
|
|
def max_per_data(self, i, onlyOneElement = False) -> Module:
|
|
mod = Module("max_per_data")
|
|
if (self.i_type.isHalf()):
|
|
mod.add(ri.VMaxF16(vgpr("Output"), vgpr("Output"), vgpr(f"Value+{i}", isAbs=True)))
|
|
# On non-Ecc hardware, the top 16 bits are dirty
|
|
if not onlyOneElement:
|
|
mod.add(ri.VLShiftRightB32(vgpr(f"Value+{i}"), 16, vgpr(f"Value+{i}")))
|
|
mod.add(ri.VMaxF16(vgpr("Output"), vgpr("Output"), vgpr(f"Value+{i}", isAbs=True)))
|
|
elif (self.i_type.isSingle()):
|
|
mod.add(ri.VMaxF32(vgpr("Output"), vgpr("Output"), vgpr(f"Value+{i}", isAbs=True)))
|
|
return mod
|
|
|
|
|
|
def scale_per_data(self, i) -> Module:
|
|
mod = Module("scale_per_data")
|
|
if self.is_scale:
|
|
mod.add(ri.VMulF32(vgpr(f"OutputD+{i}"), sgpr("Scale"), vgpr(f"Value+{i}")))
|
|
if self.scale_type == DataType("F8") or self.scale_type == DataType("F8N"):
|
|
mod.add(ri.VCmpClassF32(dst=VCC(), src0=vgpr(f"OutputD+{i}"), src1=vgpr("Fp8NanInf")))
|
|
mod.add(ri.VMed3F32(dst=vgpr("Fp8Tmp"), src0=vgpr(f"OutputD+{i}"), src1=vgpr("Fp8Min"), src2=vgpr("Fp8Max")))
|
|
mod.add(ri.VCndMaskB32(dst=vgpr(f"OutputD+{i}"), src0=vgpr("Fp8Tmp"), src1=vgpr(f"OutputD+{i}"), src2=VCC()))
|
|
mod.add(ri.VCvtPkF32toFP8(vgpr(f"OutputD+{i}"), vgpr(f"OutputD+{i}"), vgpr(f"OutputD+{i}")))
|
|
elif self.scale_type == DataType("B8") or self.scale_type == DataType("B8N"):
|
|
mod.add(ri.VCmpClassF32(dst=VCC(), src0=vgpr(f"OutputD+{i}"), src1=vgpr("BF8NanInf")))
|
|
mod.add(ri.VMed3F32(dst=vgpr("BF8Tmp"), src0=vgpr(f"OutputD+{i}"), src1=vgpr("BF8Min"), src2=vgpr("BF8Max")))
|
|
mod.add(ri.VCndMaskB32(dst=vgpr(f"OutputD+{i}"), src0=vgpr("BF8Tmp"), src1=vgpr(f"OutputD+{i}"), src2=VCC()))
|
|
mod.add(ri.VCvtPkF32toBF8(vgpr(f"OutputD+{i}"), vgpr(f"OutputD+{i}"), vgpr(f"OutputD+{i}")))
|
|
return mod
|
|
|
|
|
|
def sum_per_threadxN(self) -> Module:
|
|
mod = Module("sum_per_threadxN")
|
|
mod.addComment0("sum_per_threadxN")
|
|
mod.add(ri.SLShiftRightB32(sgpr("MainLoop"), \
|
|
int(log2(self.num_workitems * self.num_load_count * self.num_load_size * (4 // self.bpe))), \
|
|
sgpr("SizeLength")))
|
|
with asm_loop(mod, "sum_per_threadxN", "MainLoop"):
|
|
for i in range(0, self.num_load_count): # unroll
|
|
mod.add(ri.BufferLoadB128(vgpr(f"Value+{i*self.num_load_size}",4), \
|
|
vgpr(f"Offset+{i}"), sgpr("Src",4), 0, MUBUFModifiers(offen=True)))
|
|
mod.addSpaceLine()
|
|
# max operation
|
|
for i in range(0, self.num_load_count): # unroll
|
|
mod.add(ri.SWaitCnt(vlcnt=(self.num_load_count-i-1)))
|
|
for j in range(0, self.num_load_size): # dwordx4
|
|
mod.add(self.max_per_data(i * self.num_load_size + j))
|
|
mod.addSpaceLine()
|
|
# scale operation
|
|
for i in range(0, self.num_load_count): # unroll
|
|
for j in range(0, self.num_load_size): # dwordx4
|
|
mod.add(self.scale_per_data(i * self.num_load_size + j))
|
|
mod.addSpaceLine()
|
|
if self.is_scale: # buffer store fp8
|
|
for i in range(0, self.num_load_count): # unroll
|
|
for j in range(0, self.num_load_size): # dwordx4
|
|
mod.add(ri.BufferStoreB8(vgpr(f"OutputD+{i*self.num_load_size+j}"), \
|
|
vgpr(f"OffsetD+{i*self.num_load_size+j}"), \
|
|
sgpr("DstD",4), 0, MUBUFModifiers(offen=True)))
|
|
mod.addSpaceLine()
|
|
# adjust offset of buffer load
|
|
# total bytes = num_workitems * num_unroll * load_size_in_bytes
|
|
# num_unroll = num_load_count
|
|
# load_size_in_bytes = dwordx4 = num_load_size * 4
|
|
mod.add(ri.SMovB32(sgpr("Tmp"), self.num_workitems * self.num_load_count * self.num_load_size * 4))
|
|
for i in range(0, self.num_load_count):
|
|
mod.add(ri.VAddU32(vgpr(f"Offset+{i}"), vgpr(f"Offset+{i}"), sgpr("Tmp")))
|
|
mod.addSpaceLine()
|
|
if self.is_scale: # adjust offset of buffer store fp8
|
|
mod.add(ri.SMovB32(sgpr("TmpD"), self.num_workitems * self.num_load_count * 4))
|
|
for i in range(0, self.num_load_count):
|
|
for j in range(0, self.num_load_size):
|
|
mod.add(ri.VAddU32(vgpr(f"OffsetD+{i*self.num_load_size+j}"), \
|
|
vgpr(f"OffsetD+{i*self.num_load_size+j}"), sgpr("TmpD")))
|
|
mod.addSpaceLine()
|
|
mod.addSpaceLine()
|
|
return mod
|
|
|
|
|
|
def sum_per_threadx4(self) -> Module:
|
|
mod = Module("sum_per_threadx4")
|
|
mod.addComment0("sum_per_threadx4")
|
|
mod.add(ri.SLShiftRightB32(sgpr("MainLoop"), \
|
|
int(log2(self.num_workitems * self.num_load_size * (4 // self.bpe))), \
|
|
sgpr("SizeLength")))
|
|
mod.add(ri.SAndB32(sgpr("MainLoop"), hex(self.num_load_count-1), sgpr("MainLoop")))
|
|
with asm_loop(mod, "sum_per_threadx4", "MainLoop"):
|
|
mod.add(ri.BufferLoadB128(vgpr("Value",4), vgpr("Offset"), sgpr("Src",4), 0, MUBUFModifiers(offen=True)))
|
|
mod.addSpaceLine()
|
|
mod.add(ri.SWaitCnt(vlcnt=0))
|
|
# max operation
|
|
for i in range(0, self.num_load_size): # dwordx4
|
|
mod.add(self.max_per_data(i))
|
|
mod.addSpaceLine()
|
|
# scale operation
|
|
for i in range(0, self.num_load_size): # dwordx4
|
|
mod.add(self.scale_per_data(i))
|
|
mod.addSpaceLine()
|
|
if self.is_scale: # buffer store fp8
|
|
for i in range(0, self.num_load_size): # dwordx4
|
|
mod.add(ri.BufferStoreB8(vgpr(f"OutputD+{i}"), \
|
|
vgpr(f"OffsetD+{i}"), \
|
|
sgpr("DstD",4), 0, MUBUFModifiers(offen=True)))
|
|
mod.addSpaceLine()
|
|
# adjust offset of buffer load
|
|
mod.add(ri.SMovB32(sgpr("Tmp"), self.num_workitems * self.num_load_size * 4))
|
|
mod.add(ri.VAddU32(vgpr("Offset"), vgpr("Offset"), sgpr("Tmp")))
|
|
if self.is_scale: # adjust offset of buffer store fp8
|
|
mod.add(ri.SMovB32(sgpr("TmpD"), self.num_workitems * 4))
|
|
for i in range(0, self.num_load_size):
|
|
mod.add(ri.VAddU32(vgpr(f"OffsetD+{i}"), \
|
|
vgpr(f"OffsetD+{i}"), sgpr("TmpD")))
|
|
mod.addSpaceLine()
|
|
mod.addSpaceLine()
|
|
return mod
|
|
|
|
|
|
def adjust_global_address(self) -> Module:
|
|
mod = Module("adjust_global_address")
|
|
mod.addComment0("adjust_global_address")
|
|
|
|
# adjust buffer load offset
|
|
# buffer_load_dwordx4 = byte * 4 * 4
|
|
# -) buffer_load_dword = byte * 1 * 4
|
|
# --------------------------------------
|
|
# byte * 3 * 4
|
|
mod.add(ri.VMulLOU32(vgpr("Tmp"), 3, vgpr("Serial")))
|
|
mod.add(ri.VLShiftLeftB32(vgpr("Tmp"), 2, vgpr("Tmp")))
|
|
mod.add(ri.VSubU32(vgpr("Offset"), vgpr("Offset"), vgpr("Tmp")))
|
|
if (self.i_type.isHalf()):
|
|
mod.add(ri.VLShiftLeftB32(vgpr("Tmp"), 1, vgpr("Serial")))
|
|
mod.add(ri.VSubU32(vgpr("Offset"), vgpr("Offset"), vgpr("Tmp")))
|
|
mod.addSpaceLine()
|
|
|
|
if self.is_scale: # adjust buffer store offset
|
|
mod.add(ri.VMulLOU32(vgpr("TmpD"), 3, vgpr("Serial")))
|
|
mod.add(ri.VSubU32(vgpr("OffsetD"), vgpr("OffsetD"), vgpr("TmpD")))
|
|
|
|
mod.addSpaceLine()
|
|
mod.addSpaceLine()
|
|
return mod
|
|
|
|
|
|
def sum_per_thread(self) -> Module:
|
|
offset = self.num_workitems
|
|
mod = Module("sum_per_thread")
|
|
mod.addComment0("sum_per_thread")
|
|
mod.add(ri.SLShiftRightB32(sgpr("MainLoop"), int(log2(offset)), sgpr("SizeLength")))
|
|
mod.add(ri.SAndB32(sgpr("MainLoop"), sgpr("MainLoop"), self.num_load_size * (4 // self.bpe) - 1))
|
|
with asm_loop(mod, "sum_per_thread", "MainLoop"):
|
|
BufferLoadx1 = self.global_read_inst_type(1)
|
|
mod.add(BufferLoadx1(vgpr("Value"), vgpr("Offset"), sgpr("Src",4), 0, MUBUFModifiers(offen=True)))
|
|
mod.add(ri.SWaitCnt(vlcnt=0))
|
|
mod.addSpaceLine()
|
|
mod.add(self.max_per_data(0, onlyOneElement=True))
|
|
mod.addSpaceLine()
|
|
mod.add(self.scale_per_data(0))
|
|
mod.addSpaceLine()
|
|
if self.is_scale:
|
|
mod.add(ri.BufferStoreB8(vgpr("OutputD"),
|
|
vgpr("OffsetD"), sgpr("DstD",4), 0, MUBUFModifiers(offen=True)))
|
|
mod.addSpaceLine()
|
|
mod.add(ri.SMovB32(sgpr("Tmp"), self.num_workitems * self.bpe))
|
|
mod.add(ri.VAddU32(vgpr("Offset"), vgpr("Offset"), sgpr("Tmp")))
|
|
mod.addSpaceLine()
|
|
if self.is_scale:
|
|
mod.add(ri.SMovB32(sgpr("TmpD"), self.num_workitems))
|
|
mod.add(ri.VAddU32(vgpr("OffsetD"), vgpr("OffsetD"), sgpr("TmpD")))
|
|
mod.addSpaceLine()
|
|
mod.addSpaceLine()
|
|
return mod
|
|
|
|
|
|
def sum_in_some_thread(self) -> Module:
|
|
label_sum_end = Label("sum", f'loop sum end')
|
|
mod = Module("sum_in_some_thread")
|
|
mod.addComment0("sum_in_some_thread")
|
|
mod.add(ri.SAndB32(sgpr("MainLoop"), sgpr("SizeLength"), self.num_workitems-1))
|
|
mod.add(ri.VCmpLtU32(VCC(), vgpr("Serial"), sgpr("MainLoop")))
|
|
mod.add(ri.SCBranchVCCZ(label_sum_end.getLabelName()))
|
|
mod.add(self.SMovBX(EXEC(), VCC()))
|
|
mod.add(ri.SNop(1))
|
|
BufferLoadx1 = self.global_read_inst_type(1)
|
|
mod.add(BufferLoadx1(vgpr("Value"), vgpr("Offset"), sgpr("Src",4), 0, MUBUFModifiers(offen=True)))
|
|
mod.add(ri.SWaitCnt(vlcnt=0))
|
|
mod.addSpaceLine()
|
|
mod.add(self.max_per_data(0, onlyOneElement=True))
|
|
mod.addSpaceLine()
|
|
mod.add(self.scale_per_data(0))
|
|
mod.addSpaceLine()
|
|
if self.is_scale:
|
|
mod.add(ri.BufferStoreB8(vgpr("OutputD"),
|
|
vgpr("OffsetD"), sgpr("DstD",4), 0, MUBUFModifiers(offen=True)))
|
|
mod.addSpaceLine()
|
|
mod.add(self.SMovBX(EXEC(), "-1"))
|
|
mod.add(ri.SNop(1))
|
|
mod.add(label_sum_end)
|
|
mod.addSpaceLine()
|
|
mod.addSpaceLine()
|
|
return mod
|
|
|
|
|
|
def merge_sum(self) -> Module:
|
|
mod = Module("merge_sum")
|
|
if (self.i_type.isHalf()):
|
|
mod.add(ri.VMaxF16(vgpr("Output"), vgpr("Output"), vgpr("OutputB")))
|
|
elif (self.i_type.isSingle()):
|
|
mod.add(ri.VMaxF32(vgpr("Output"), vgpr("Output"), vgpr("OutputB")))
|
|
|
|
return mod
|
|
|
|
|
|
def intra_wave_reduction(self) -> Module:
|
|
label = Label("permute", "permute")
|
|
mod = Module("intra_wave_reduction")
|
|
mod.addComment0("intra_wave_reduction")
|
|
mod.add(ri.SMovB32(sgpr("Tmp"), 1))
|
|
mod.add(label)
|
|
mod.addSpaceLine()
|
|
mod.add(ri.VAddU32(vgpr("Tmp"), sgpr("Tmp"), vgpr("Serial")))
|
|
mod.add(ri.VAndB32(vgpr("Tmp"), self.wavefront_size-1, vgpr("Tmp")))
|
|
mod.add(ri.VLShiftLeftB32(vgpr("Tmp"), 0x2, vgpr("Tmp")))
|
|
mod.addSpaceLine()
|
|
mod.add(ri.DSBPermuteB32(vgpr("OutputB"), vgpr("Tmp"), vgpr("Output")))
|
|
mod.add(ri.SWaitCnt(dscnt=0))
|
|
mod.addSpaceLine()
|
|
mod.add(self.merge_sum())
|
|
mod.add(ri.SLShiftLeftB32(sgpr("Tmp"), 1, sgpr("Tmp")))
|
|
mod.add(ri.SCmpLtU32(sgpr("Tmp"), self.wavefront_size))
|
|
mod.add(ri.SCBranchSCC1(label.getLabelName()))
|
|
mod.addSpaceLine()
|
|
return mod
|
|
|
|
|
|
def inter_wave_reduction(self) -> Module:
|
|
label_inter = Label("inter", f'inter')
|
|
label_upper = Label("upper", f'upper')
|
|
label_lower = Label("lower", f'lower')
|
|
label_empty = Label("empty", f'empty')
|
|
label_end = Label("end", f'end')
|
|
mod = Module("inter_wave_reduction")
|
|
mod.addComment0("inter_wave_reduction")
|
|
mod.add(ri.VLShiftRightB32(vgpr("Widx"), int(log2(self.wavefront_size)), vgpr("Serial")))
|
|
mod.add(ri.SMovB32(sgpr("Offset"), self.num_workitems // self.wavefront_size))
|
|
mod.add(label_inter)
|
|
mod.add(ri.SLShiftRightB32(sgpr("Offset"), 1, sgpr("Offset")))
|
|
mod.add(ri.SCmpEQU32(sgpr("Offset"), 0))
|
|
mod.add(ri.SCBranchSCC1(label_end.getLabelName()))
|
|
mod.add(ri.SLShiftLeftB32(sgpr("Tmp"), 1, sgpr("Offset")))
|
|
mod.add(ri.VCmpLtU32(sgpr("Tmp+2",self.laneSGPRCount), vgpr("Widx"), sgpr("Tmp")))
|
|
mod.add(ri.VCmpGEU32(sgpr("Tmp+4",self.laneSGPRCount), vgpr("Widx"), sgpr("Offset")))
|
|
mod.add(self.SAndBX(VCC(), sgpr("Tmp+2",self.laneSGPRCount), sgpr("Tmp+4",self.laneSGPRCount)))
|
|
mod.add(ri.SCBranchVCCNZ(label_upper.getLabelName()))
|
|
mod.add(ri.VCmpLtU32(VCC(), vgpr("Widx"), sgpr("Offset")))
|
|
mod.add(ri.SCBranchVCCNZ(label_lower.getLabelName()))
|
|
mod.add(ri.SBranch(label_empty.getLabelName()))
|
|
|
|
mod.add(label_upper)
|
|
mod.add(ri.VSubU32(vgpr("Tmp"), vgpr("Widx"), sgpr("Offset")))
|
|
mod.add(ri.VLShiftLeftB32(vgpr("Tmp"), int(log2(self.bpe)), vgpr("Tmp")))
|
|
ds = DSModifiers(offset=0)
|
|
DSStorex1 = self.local_write_inst_type(1)
|
|
mod.add(DSStorex1(vgpr("Tmp"), vgpr("Output"), ds))
|
|
mod.add(ri.SWaitCnt(dscnt=0))
|
|
mod.add(ri.SBarrier())
|
|
mod.add(ri.SBranch(label_inter.getLabelName()))
|
|
mod.add(label_lower)
|
|
mod.add(ri.SBarrier())
|
|
mod.add(ri.VLShiftLeftB32(vgpr("Tmp"), int(log2(self.bpe)), vgpr("Widx")))
|
|
ds = DSModifiers(offset=0)
|
|
DSLoadx1 = self.local_read_inst_type(1)
|
|
mod.add(DSLoadx1(vgpr("OutputB"), vgpr("Tmp"), ds))
|
|
mod.add(ri.SWaitCnt(dscnt=0))
|
|
mod.add(self.merge_sum())
|
|
mod.add(ri.SBranch(label_inter.getLabelName()))
|
|
mod.add(label_empty)
|
|
mod.add(ri.SBarrier())
|
|
mod.add(ri.SBranch(label_inter.getLabelName()))
|
|
mod.add(label_end)
|
|
mod.addSpaceLine()
|
|
return mod
|
|
|
|
|
|
def broadcast(self) -> Module:
|
|
label_lower = Label("broadcast_lower", f'broadcast_lower')
|
|
label_end = Label("broadcast_end", f'broadcast_end')
|
|
|
|
mod = Module("broadcast")
|
|
mod.addComment0("broadcast")
|
|
mod.add(ri.VCmpEQU32(VCC(), vgpr("Widx"), 0))
|
|
mod.add(ri.SCBranchVCCZ(label_lower.getLabelName()))
|
|
ds = DSModifiers(offset=0)
|
|
DSStorex1 = self.local_write_inst_type(1)
|
|
mod.add(DSStorex1(vgpr("Widx"), vgpr("Output"), ds))
|
|
mod.add(ri.SWaitCnt(dscnt=0))
|
|
mod.add(ri.SBarrier())
|
|
mod.add(ri.SBranch(label_end.getLabelName()))
|
|
mod.add(label_lower)
|
|
mod.add(ri.SBarrier())
|
|
mod.add(ri.VMovB32(vgpr("Tmp"), 0))
|
|
ds = DSModifiers(offset=0)
|
|
DSLoadx1 = self.local_read_inst_type(1)
|
|
mod.add(DSLoadx1(vgpr("Output"), vgpr("Tmp"), ds))
|
|
mod.add(ri.SWaitCnt(dscnt=0))
|
|
mod.add(label_end)
|
|
mod.addSpaceLine()
|
|
return mod
|
|
|
|
|
|
def output_result(self) -> Module:
|
|
mod = Module("output_result")
|
|
mod.addComment0("output_result")
|
|
BufferStorex1 = self.global_write_inst_type(1)
|
|
|
|
mod.add(ri.VMovB32(vgpr("Offset"), 0))
|
|
if self.i_type.toChar() == 'H' and self.o_type.toChar() == "S":
|
|
mod.add(ri.VCvtF16toF32(vgpr("Output"), vgpr("Output")))
|
|
elif self.i_type.toChar() == 'S' and self.o_type.toChar() == "H":
|
|
mod.add(ri.VCvtF32toF16(vgpr("Output"), vgpr("Output")))
|
|
mod.add(BufferStorex1(vgpr("Output"), vgpr("Offset"), sgpr("Dst",4), 0, MUBUFModifiers(offen=True)))
|
|
mod.addSpaceLine()
|
|
|
|
return mod
|
|
|
|
def amax_kernel_body(self) -> Module:
|
|
mod = Module(self.func_name)
|
|
mod.add(self.defineVariables())
|
|
with asm_func(self.func_name, mod):
|
|
mod.add(self.load_kernel_args())
|
|
mod.add(self.init_param())
|
|
mod.add(self.calculate_global_address())
|
|
mod.add(self.sum_per_threadxN())
|
|
mod.add(self.sum_per_threadx4())
|
|
mod.add(self.adjust_global_address())
|
|
mod.add(self.sum_per_thread())
|
|
mod.add(self.sum_in_some_thread())
|
|
mod.add(self.intra_wave_reduction())
|
|
mod.add(self.inter_wave_reduction())
|
|
mod.add(self.broadcast())
|
|
mod.add(self.output_result())
|
|
return mod
|
|
|
|
@dataclass
|
|
class KernelArgument:
|
|
size: int
|
|
offset: int
|
|
value_kind: str
|
|
address_space: Optional[str] = None
|
|
|
|
def to_dict(self):
|
|
d = {'.size': self.size, '.offset': self.offset,
|
|
'.value_kind': self.value_kind}
|
|
|
|
if self.address_space:
|
|
d['.address_space'] = self.address_space
|
|
|
|
return d
|
|
|
|
@dataclass
|
|
class KernelMeta:
|
|
name: str
|
|
num_vgpr: int
|
|
num_sgpr: int
|
|
num_agpr: int
|
|
num_lds_bytes: int
|
|
wavefront_size: int
|
|
max_workgroup_size: int
|
|
args_alignment: int
|
|
args: List[KernelArgument]
|
|
|
|
def update_args_offsets(self):
|
|
offset = 0
|
|
for arg in args:
|
|
arg.offset = offset
|
|
offset += arg.size
|
|
|
|
def _get_args_size(self):
|
|
total_size = sum(arg.size for arg in self.args)
|
|
total_size += (self.args_alignment - (total_size % self.args_alignment))
|
|
return total_size
|
|
|
|
def to_dict(self):
|
|
return {
|
|
'.name': self.name,
|
|
'.symbol': f'{self.name}.kd',
|
|
'.kernarg_segment_size': self._get_args_size(),
|
|
'.group_segment_fixed_size': self.num_lds_bytes,
|
|
'.private_segment_fixed_size': 0,
|
|
'.kernarg_segment_align': self.args_alignment,
|
|
'.wavefront_size': self.wavefront_size,
|
|
'.sgpr_count': self.num_sgpr,
|
|
'.vgpr_count': self.num_vgpr,
|
|
'.agpr_count': self.num_agpr,
|
|
'.max_flat_workgroup_size': self.max_workgroup_size,
|
|
'.args': [arg.to_dict() for arg in self.args]
|
|
}
|
|
|
|
def meta_str(kernels: Tuple[KernelMeta]):
|
|
beg = '.amdgpu_metadata\n---'
|
|
content_str = yaml.dump({'amdhsa.version': [1, 1], 'amdhsa.kernels': [kernel.to_dict() for kernel in kernels]})
|
|
end = '.end_amdgpu_metadata'
|
|
return '\n'.join([beg, content_str, end, ''])
|
|
|
|
|
|
if __name__ == '__main__':
|
|
ap = ArgumentParser()
|
|
ap.add_argument('-o', '--output', type=str, required=True, help='Output path of compiled binary')
|
|
ap.add_argument('-t', type=str, default="S", help='data type')
|
|
ap.add_argument('-d', type=str, default="None", help='dest data type')
|
|
ap.add_argument('-s', type=str, default="F8N", help='scale data type')
|
|
ap.add_argument('-w', type=int, default=256, help='workitem')
|
|
ap.add_argument('-c', type=int, default=4, help='load conut per iteration')
|
|
ap.add_argument('--toolchain', type=str, default=ToolchainDefaults.CXX_COMPILER, help='Path to ROCm compiler')
|
|
ap.add_argument('--debug-build', action='store_true', dest='debug_build', help='Build with debug information')
|
|
ap.add_argument('--is-scale', action='store_true', dest='is_scale', help='Enable scaled output or not')
|
|
ap.add_argument('--arch', type=str, default='gfx90a', help='Target architecture for assembler, e.g. gfx908. Default is gfx90a')
|
|
ap.add_argument('--xnack', action='store_true', help='Append :xnack+ to the .amdgcn_target code-object id (arch logic still uses the base arch)')
|
|
ap.set_defaults(debug_build=False)
|
|
|
|
args = ap.parse_args()
|
|
output_path: str = args.output
|
|
t: str = args.t
|
|
d: str = t if (args.d =="None") else args.d
|
|
s: str = args.s
|
|
w: int = args.w
|
|
c: int = args.c
|
|
toolchain_path: str = validateToolchain(args.toolchain)
|
|
debug_build: bool = args.debug_build
|
|
arch: str = args.arch
|
|
is_scale: bool = args.is_scale
|
|
xnack: bool = args.xnack
|
|
isa = gfxToIsa(arch)
|
|
|
|
if any([not i for i in (arch, toolchain_path, isa)]):
|
|
restoreDefaultGlobalParameters()
|
|
assignGlobalParameters({})
|
|
enumerator = validateToolchain(ToolchainDefaults.DEVICE_ENUMERATOR)
|
|
isa = detectGlobalCurrentISA(0, enumerator)
|
|
arch = isaToGfx(isa)
|
|
toolchain_path = validateToolchain(ToolchainDefaults.CXX_COMPILER)
|
|
|
|
_global_ti.init(isa, toolchain_path, False)
|
|
waveFrontSize = 32 if isa[0] in [11, 12] else 64
|
|
_global_ti.setKernel(isa, waveFrontSize)
|
|
amax = AMaxKernelGenerator(DataType(t), DataType(d), DataType(s), w, c, 4, waveFrontSize, arch, isa, is_scale)
|
|
kernel_body = amax.amax_kernel_body()
|
|
args = amax.kernel_args()
|
|
func_name = amax.func_name
|
|
meta = KernelMeta(func_name, amax.vgpr_pool.size(), amax.sgpr_pool.size(), 0, amax.lds_usage_byte, waveFrontSize, w, 8, args)
|
|
meta.update_args_offsets()
|
|
k_str = '\n'.join([kernel_header(func_name, arch, amax.vgpr_pool.size(), amax.sgpr_pool.size(), amax.lds_usage_byte, xnack),
|
|
meta_str((meta,)),
|
|
str(kernel_body)])
|
|
|
|
with open(output_path, 'w') as f:
|
|
f.write(k_str)
|
|
|
|
output_path_basename = os.path.splitext(output_path)[0]
|
|
amax.dump('yaml', f'{output_path_basename}.yaml')
|