KGCompiler: Deep Learning Compilation Optimization for Knowledge Graph Complex Logical Query Answering

TL;DR

KGCompiler presents a KG-oriented deep learning compiler for Complex Logical Query Answering on Knowledge Graphs. By converting FOL queries into computation graphs, identifying operator patterns, and applying horizontal, vertical, and hybrid fusion, it achieves substantial speedups (up to 8.26x, average ~3.71x) and memory reductions without compromising accuracy. The approach is implemented on TorchInductor and validated on 14 CLQA tasks across standard KG datasets, demonstrating broad applicability and efficiency gains for practical KG reasoning systems.

Abstract

Complex Logical Query Answering (CLQA) involves intricate multi-hop logical reasoning over large-scale and potentially incomplete Knowledge Graphs (KGs). Although existing CLQA algorithms achieve high accuracy in answering such queries, their reasoning time and memory usage scale significantly with the number of First-Order Logic (FOL) operators involved, creating serious challenges for practical deployment. In addition, current research primarily focuses on algorithm-level optimizations for CLQA tasks, often overlooking compiler-level optimizations, which can offer greater generality and scalability. To address these limitations, we introduce a Knowledge Graph Compiler, namely KGCompiler, the first deep learning compiler specifically designed for CLQA tasks. By incorporating KG-specific optimizations proposed in this paper, KGCompiler enhances the reasoning performance of CLQA algorithms without requiring additional manual modifications to their implementations. At the same time, it significantly reduces memory usage. Extensive experiments demonstrate that KGCompiler accelerates CLQA algorithms by factors ranging from 1.04x to 8.26x, with an average speedup of 3.71x. We also provide an interface to enable hands-on experience with KGCompiler.

Figures (7)

• Figure 1: An illustration of the different types of CLQA tasks. Symbol explanation: $\land$ represents intersection, $\lor$ represents union, $\neg$ represents negation, and $\rightarrow$ represents projection. FOL operators naming convention: $p$ - projection, $i$ - intersection, $n$ - negation, $u$ - union.
• Figure 2: Overview of the FOL query process mapped to the Computation Graph: The complex Logical Query (A) is first parsed into a Query Structure (B). Subsequently, the query Computation Plan (C) is followed to execute the query directly within the Computation Graph (D).
• Figure 3: Architecture of KGCompiler. To begin with, the Logical Query (A) is parsed into a Query Structure (B), which is then mapped to a Computation Plan (C). Next, Graph Capturer processes this Computation Plan to construct the Original Computation Graph (D). Following this, Pattern Recognizer extracts the combination patterns in the graph, resulting in a Modularized Computation Graph (E). Finally, by applying the fusion optimizations based on the recognized patterns, Operator Fuser generates the Optimized Computation Graph (F).
• Figure 4: Take the BetaE algorithm as an example, under each task type and different batch sizes, the time of processing a single query using KGReasoning (blue line) and KGCompiler (orange line) is measured respectively.
• Figure 5: Take the BetaE algorithm as an example, the maximum memory footprint during computation is measured for KGReasoning (blue line) and KGCompiler (orange line) under each task type and different batch sizes.