pack_featgraph.py

""" Export featgraph kernels to a shared library. """
import tvm
from sddmm import sddmm_tree_reduction_gpu


def get_sddmm_kernels_gpu(idtypes, dtypes):
    """
    Parameters
    ----------
    idtypes: List[str]
        Possible index types.
    dtypes: List[str]
        Possible data types.

    Returns
    -------
    List[IRModule]:
        The list of IRModules.
    """
    ret = []
    # SDDMM Tree Reduction
    for dtype in dtypes:
        for idtype in idtypes:
            ret.append(sddmm_tree_reduction_gpu(idtype, dtype))

    return ret


if __name__ == '__main__':
    binary_path = 'libfeatgraph_kernels.so'
    kernels = []
    idtypes = ['int32', 'int64']
    dtypes = ['float16', 'float64', 'float32', 'int32', 'int64']

    kernels += get_sddmm_kernels_gpu(idtypes, dtypes)

    # build kernels and export the module to libfeatgraph_kernels.so
    module = tvm.build(kernels, target='cuda', target_host='llvm')
    module.export_library(binary_path)