Conference proceeding
cuKE: An Efficient Code Generator for Score Function Computation in Knowledge Graph Embedding
2024 IEEE International Parallel and Distributed Processing Symposium (IPDPS), pp.903-914
05/27/2024
DOI: 10.1109/IPDPS57955.2024.00085
Abstract
Knowledge graph embedding (KGE) plays an important role in graph mining and learning applications by converting discrete graph structures to continuous vector representations. While previous systems have focused on scaling KGE onto multiple GPUs, the score function computation on each GPU can be a performance bottleneck. Existing KGE systems implement the score functions with separate tensor operations, leading to large memory consumption and poor memory access efficiency. To overcome the issues, we propose a code generator that automatically translates Python-like definitions of KGE score functions into efficient CUDA code. Our code generator exploits the unique feature of KGE score functions and performs an aggressive fusion of tensor operations. Additionally, our generated code performs a runtime inspection to reduce redundant memory access for edges with identical indices. Experiments show that our generated code uses much less memory than previous systems and achieves an average speedup of 14.9x over TorchScript and 7.8x over TVM.
Details
- Title: Subtitle
- cuKE: An Efficient Code Generator for Score Function Computation in Knowledge Graph Embedding
- Creators
- Lihan Hu - University of IowaJing Li - NvidiaPeng Jiang - University of Iowa
- Resource Type
- Conference proceeding
- Publication Details
- 2024 IEEE International Parallel and Distributed Processing Symposium (IPDPS), pp.903-914
- Publisher
- IEEE
- DOI
- 10.1109/IPDPS57955.2024.00085
- ISSN
- 1530-2075
- eISSN
- 1530-2075
- Grant note
- NSF: CNS-2310423
This work was supported by NSF award CNS-2310423.
- Language
- English
- Date published
- 05/27/2024
- Academic Unit
- Computer Science
- Record Identifier
- 9984658348502771
Metrics
2 Record Views