Way too much overhead for ndata and edata
- 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.