dglai / featgraph Goto Github PK

• Fix the dlpack bug: we found that we'll meet segfault irregularly when integrating featgraph to DGL.
  • TODO(zihao): find a minimal example to reproduce the error.
• Support broadcasting
  • TODO(zihao): investigate how to support general broadcasting.
• SPMM with edge weight
• Process node/edge mapping on DGL end before feeding the tensors to featgraph.

The kernels are compiled for flatten feature, i.e., has two dimensions: the first is number of nodes/edges, the second is flatten feature length. For example, feature shape (2,3) will become (6,).

The problem occurs when unflatten feature is given to the kernel as input. There are several situations:

# if you use one dimension feature, no problem
# lhs_shp = (6,)
# but if you use more than one dimension, it is problematic
lhs_shp = (2,3)

import torch as th
# if you import tvm before you import dgl.sparse, you will get assert fail
import tvm
# if you does not import tvm, or import it after dgl.sparse, then kernel will segfault
import dgl.sparse

num_nodes = 10
num_edges = 5
# by default use int64 as indice type. same behavior for int32 and int64
g = dgl.rand_graph(num_nodes, num_edges)
gidx = g._graph

u = th.rand((num_nodes,)+lhs_shp, dtype=th.float32, device=th.device('cpu'))
# use copy_lhs, so broadcasting not related
# but it is the same if you use other op, or even use broadcast
out = dgl.sparse._gspmm(gidx, 'copy_lhs', 'sum', u, None)
print(out)

To reproduce the bugs, first you need to use the integrated kernel, and this featgraph repo on lnzhi-dev branch. Then inside the sparse.py, you need to comment out two lines, 226 and 229. Now you should see the bugs I describe.

Finally, I think we actually do need to reshape the feature to be same as we define in the kernel, as the assertion failure says. But whether you import tvm decides whether you get an assertion error or segmentation fault is quite interesting, as we do not use anything from tvm here.

@yzh119 @Huyuwei