dmlc / dmlc/dgl

Way too much overhead for ndata and edata

Open
#3,317 2 comments 0 reactions 0 assignees View on GitHub
feature request topic: system performance
Dominant language
Python
Stars
14.3k
Forks
3.1k
PR merge metrics
No merged PRs in 30d

Description

During our optimization of RGCN we found that the following interfaces have non-negligible Python overhead to the extent that the optimization in heterogeneous graph kernel did not bring much benefit.
* `ndata`, `edata`, `nodes`, `edges`
* `apply_edges`
* `update_all`

For instance, the following two styles will give 0.6ms and 0.1ms running time respectively with the same exact arguments (credit to @isratnisa, maybe you can give more details on profiling and correct the numbers if I'm wrong?):
```python
g.edges[rel].data['h*w_r'] = th.matmul(g.edges[rel].data['m'], weight[rel]) # 0.6ms
h_t[rel] = th.matmul(h_t[rel], weight[rel]) # 0.1ms
```

Some initial profiling on `ndata` showed that `ndata['x']` is 200x slower than a dictionary update:
```python
import torch
import dgl
import time

M = 1000
N = 10000

g = dgl.graph((torch.randint(0, M, (N,)), torch.randint(0, M, (N,))))
x = torch.randn(M, 200)
d = {}

t0 = time.time()
for _ in range(1000000):
g.ndata['x'] = x
tt = time.time()
print(tt - t0) # 17.0s

t0 = time.time()
for _ in range(1000000):
d['x'] = x
tt = time.time()
print(tt - t0) # 0.083s
```

The majority of time is spent in type checking and unnecessary Python calls, as shown below.
```
Timer unit: 1e-06 s

Total time: 14.6243 s
File: /mnt/b/gq/dgl/python/dgl/frame.py
Function: update_column at line 497

Line # Hits Time Per Hit % Time Line Contents
==============================================================
497 @profile
498 def update_column(self, name, data):
499 """Add or replace the column with the given name and data.
500
501 Parameters
502 ----------
503 name : str
504 The column name.
505 data : Column or data convertible to Column
506 The column data.
507 """
508 1000000 8877575.0 8.9 60.7 col = Column.create(data)
509 1000000 4071697.0 4.1 27.8 if len(col) != self.num_rows:
510 raise DGLError('Expected data to have %d rows, got %d.' %
511 (self.num_rows, len(col)))
512 1000000 1675047.0 1.7 11.5 self._columns[name] = col

Total time: 95.2076 s
File: /mnt/b/gq/dgl/python/dgl/heterograph.py
Function: _set_n_repr at line 4083

Line # Hits Time Per Hit % Time Line Contents
==============================================================
4083 @profile
4084 def _set_n_repr(self, ntid, u, data):
4085 """Internal API to set node features.
4086
4087 `data` is a dictionary from the feature name to feature tensor. Each tensor
4088 is of shape (B, D1, D2, ...), where B is the number of nodes to be updated,
4089 and (D1, D2, ...) be the shape of the node representation tensor. The
4090 length of the given node ids must match B (i.e, len(u) == B).
4091
4092 All updates will be done out of place to work with autograd.
4093
4094 Parameters
4095 ----------
4096 ntid : int
4097 Node type id.
4098 u : node, container or tensor
4099 The node(s).
4100 data : dict of tensor
4101 Node representation.
4102 """
4103 1000000 1876796.0 1.9 2.0 if is_all(u):
4104 1000000 28771379.0 28.8 30.2 num_nodes = self._graph.number_of_nodes(ntid)
4105 else:
4106 u = utils.prepare_tensor(self, u, 'u')
4107 num_nodes = len(u)
4108 2000000 3105292.0 1.6 3.3 for key, val in data.items():
4109 1000000 2708365.0 2.7 2.8 nfeats = F.shape(val)[0]
4110 1000000 1247508.0 1.2 1.3 if nfeats != num_nodes:
4111 raise DGLError('Expect number of features to match number of nodes (len(u)).'
4112 ' Got %d and %d instead.' % (nfeats, num_nodes))
4113 1000000 25528712.0 25.5 26.8 if F.context(val) != self.device:
4114 raise DGLError('Cannot assign node feature "{}" on device {} to a graph on'
4115 ' device {}. Call DGLGraph.to() to copy the graph to the'
4116 ' same device.'.format(key, F.context(val), self.device))
4117
4118 1000000 1941297.0 1.9 2.0 if is_all(u):
4119 1000000 30028221.0 30.0 31.5 self._node_frames[ntid].update(data)
4120 else:
4121 self._node_frames[ntid].update_row(u, data)

Total time: 114.841 s
File: /mnt/b/gq/dgl/python/dgl/view.py
Function: __setitem__ at line 68

Line # Hits Time Per Hit % Time Line Contents
==============================================================
68 @profile
69 def __setitem__(self, key, val):
70 1000000 1547822.0 1.5 1.3 if isinstance(self._ntype, list):
71 assert isinstance(val, dict), \
72 'Current HeteroNodeDataView has multiple node types, ' \
73 'please passing the node type and the corresponding data through a dict.'
74
75 for (ntype, data) in val.items():
76 ntid = self._graph.get_ntype_id(ntype)
77 self._graph._set_n_repr(ntid, self._nodes, {key : data})
78 else:
79 1000000 1417516.0 1.4 1.2 assert isinstance(val, dict) is False, \
80 'The HeteroNodeDataView has only one node type. ' \
81 'please pass a tensor directly'
82 1000000 111875185.0 111.9 97.4 self._graph._set_n_repr(self._ntid, self._nodes, {key : val})
```

We will gradually add more profiling results in the future.

Contributor guide

No contributing guide indexed for this repository

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.