dmlc--dgl
Hongzhi (Steve), Chen
dce899190e
* auto-reformat * lintrunner --------- Co-authored-by: Ubuntu <ubuntu@ip-172-31-28-63.ap-northeast-1.compute.internal>
40 行
975 B
Python
40 行
975 B
Python
""" 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)
|