The purpose of this prototype repo is two-fold:
- Exposing DGL's sparse matrix APIs directly to users.
- Intergrating DGL's sparse matrix ops (with autograd) into TorchScript.
In the lowest level of design, DGL uses aten::CSRMatrix
and aten::COOMatrix
to warp all the memory storage, CPU and GPU kernels. Based on them, UnitGraph
and HeteroGraph
provide higher level of abstraction to hide the concrete formats of sparse matrix and the conversion among different formats.
In the Python side, DGL implements the forward and backward versions of all the sparse ops using the pybind of HeteroGraph
, and regisers them into Pytorch as autograd.Function
. Finially, DGLGraph
uses these ops to implement the functions used by users, such as update_all
.
However, the sparse matrix APIs are not directly exposed to the users, and they can not be recognized by TorchScript, which requires the op be written in C++.
We reuse aten::CSRMatrix
and aten::COOMatrix
for low-level storage and kernels, but propose a new abstraction called AjdMatrix
responsible for managing different formats and conversion. Then we reimplement the autograd of sparse op in C++ and register them into TorchScript.