How dgl.ops.gspmm work for different message and reduce functions?

[tariqaf]

I want to understand two basic questions regarding gspmm?

• One e.g How the same gspmm performs on different message functions and reduce function operations? e.g. u_add_e_sum and u_mul_e_sum? or u_mul_e_max?

• Second thing I want to know is which C Code from src is responsible for these functions in the CPU setting? Whether they have been done by some third-party library like cusparse or libxsmm? and also which operations are solely implemented by DGL and what operations are done by 3rd party libraries (libxsmm etc) or backend like PyTorch etc? basically what I wanted to know is what percentage of DGL implementation and other libraries and even backend are involved in this different kind of gspmm operations (e.g. u_add_e_sum and u_mul_e_sum? or u_mul_e_max?). Thanks in advance for the valuable answers.

[czkkkkkk]

Hi tariqaf.

1. dgl/ops/spmm.py registers different types of message and reduce functions.
2. By tracking the definition of gspmm_internal in dgl/ops/spmm.py, you can find its C++ interface is called in dgl/sparse.py. The calling stack is dgl/ops/spmm.py → dgl/backend/pytorch/sparse.py → dgl/sparse.py

[tariqaf]

DGL community has been very helpful. Thank you so much…
Can you say a little about C code calling stack e.g for GCN with copy as message and sum as reduce function in update_all method? I wanted to know Whether DGL builtin code is called or some third party like cuSparse for GPU or Libsxmm for CPU? and what is the role of pytorch Matrix multiplication? Thank you so much in advance.

[czkkkkkk]

The calling stack of C++ code using GPU is:
src/array/kernel.cc:_CAPI_DGLKernelSpMM,
src/array/kernel.cc:SpMM(),
src/array/cuda/spmm.cu:SpMMCsr().

Depending on the message and reduce functions, DGL calls cusparse or uses its own kernels.

[tariqaf]

Thank you so much… I was having difficulty to debug C code. but this will help me understand the implementation. Thanks Again.