tilelang.contrib.rocm

Utility for ROCm backend

Functions

find_lld([required])	Find ld.lld in system.
rocm_link(in_file, out_file[, lld])	Link relocatable ELF object to shared ELF object using lld
callback_rocm_link(obj_bin)	Links object file generated from LLVM to HSA Code Object
callback_rocm_bitcode_path([rocdl_dir])	Utility function to find ROCm device library bitcodes
parse_compute_version(compute_version)	Parse compute capability string to divide major and minor version
have_matrixcore([compute_version])	Either MatrixCore support is provided in the compute capability or not
get_rocm_arch([rocm_path])	Utility function to get the AMD GPU architecture
find_rocm_path()	Utility function to find ROCm path

Module Contents

tilelang.contrib.rocm.find_lld(required=True)

Find ld.lld in system.

Parameters:

required (bool) – Whether it is required, runtime error will be raised if the compiler is required.

Returns:

valid_list – List of possible paths.

Return type:

list of str

Note

This function will first search ld.lld that matches the major llvm version that built with tvm

tilelang.contrib.rocm.rocm_link(in_file, out_file, lld=None)

Link relocatable ELF object to shared ELF object using lld

Parameters:

• in_file (str) – Input file name (relocatable ELF object file)

• out_file (str) – Output file name (shared ELF object file)

• lld (str, optional) – The lld linker, if not specified, we will try to guess the matched clang version.

tilelang.contrib.rocm.callback_rocm_link(obj_bin)

Links object file generated from LLVM to HSA Code Object

Parameters:

obj_bin (bytearray) – The object file

Returns:

cobj_bin – The HSA Code Object

Return type:

bytearray

tilelang.contrib.rocm.callback_rocm_bitcode_path(rocdl_dir=None)

Utility function to find ROCm device library bitcodes

Parameters:

rocdl_dir (str) – The path to rocm library directory The default value is the standard location

tilelang.contrib.rocm.parse_compute_version(compute_version)

Parse compute capability string to divide major and minor version

Parameters:

compute_version (str) – compute capability of a GPU (e.g. “6.0”)

Returns:

• major (int) – major version number

• minor (int) – minor version number

tilelang.contrib.rocm.have_matrixcore(compute_version=None)

Either MatrixCore support is provided in the compute capability or not

Parameters:

compute_version (str, optional) – compute capability of a GPU (e.g. “7.0”).

Returns:

have_matrixcore – True if MatrixCore support is provided, False otherwise

Return type:

bool

tilelang.contrib.rocm.get_rocm_arch(rocm_path='/opt/rocm')

Utility function to get the AMD GPU architecture

Parameters:

rocm_path (str) – The path to rocm installation directory

Returns:

gpu_arch – The AMD GPU architecture

Return type:

str

tilelang.contrib.rocm.find_rocm_path()

Utility function to find ROCm path

Returns:

path – Path to ROCm root.

Return type:

str