Deep learning compiler is becoming a fundamental piece of modern deep learning frameworks. In the domain of GNN system, there have been a rise of interest in utilizing compilers for optimizing GNNs. For DGL, we plan to integrate torch.fx and torchscript to open up these chances.
The major road blocker is the data class DGLGraph which is currently opaque to PyTorch. The internal implementation of DGLGraph has multiple layers as well as mutable members (for caching different sparse formats). Therefore, the plan is to tackle different levels of components piece by piece.
[ ] Expose the internal sparse matrix and its operators to Python using PyTorch's native FFI system.
[ ] Reimplement sparse matrix operations and their gradient operators using PyTorch's C++ extension.
[ ] Make the sparse matrix data type and operations compatible with torch.fx and torchscript.
[ ] Make DGLGraph a pure Python wrapper of the exposed sparse matrix data type.
[ ] Annotate all the exposed operators with proper type hints.
[ ] Make homogeneous DGLGraph and its operations compatible with torch.fx and torchscript.
[ ] Make heterogeneous DGLGraph and its operations compatible with torch.fx and torchscript.
[ ] Integrate compiler optimizations from Graphiler paper (MLSys'21)
[ ] Provide compiler-assisted solutions for automatic model inference.
Deep learning compiler is becoming a fundamental piece of modern deep learning frameworks. In the domain of GNN system, there have been a rise of interest in utilizing compilers for optimizing GNNs. For DGL, we plan to integrate
torch.fx
andtorchscript
to open up these chances.The major road blocker is the data class
DGLGraph
which is currently opaque to PyTorch. The internal implementation ofDGLGraph
has multiple layers as well as mutable members (for caching different sparse formats). Therefore, the plan is to tackle different levels of components piece by piece.torch.fx
andtorchscript
.DGLGraph
a pure Python wrapper of the exposed sparse matrix data type.DGLGraph
and its operations compatible withtorch.fx
andtorchscript
.DGLGraph
and its operations compatible withtorch.fx
andtorchscript
.