[Feat request] Allow multiple built-in function usage in .update_all

[felipemello1]

🚀 Feature

support multiple built-in function at once: g.update_all( [m_func1, m_func2] , [r_func1, r_func2] )

Motivation

Main motivation would be better performance, since I imagine it would eliminate considerable overhead. If there is no overhead, then the improvement would be only having a cleaner code.

Pitch

Desired solution

import dgl
import dgl.function as fn
import torch

g = dgl.heterograph({('user', 'follows', 'user'): ([0, 1, 2], [1, 2, 2])})
g.nodes['user'].data['h'] = torch.tensor([[0.], [1.], [2.]])
g.nodes['user'].data['x'] = torch.tensor([[0.], [1.], [2.]])

#CURRENT APPROACH
g['follows'].update_all(fn.copy_u('h', 'm'), fn.sum('m', 'h_new'), etype='follows')
g['follows'].update_all(fn.copy_u('x', 'm2'), fn.sum('m2', 'h_new'), etype='follows')

#DESIRED APPROACH
g['follows'].update_all([fn.copy_src('h', 'm'), fn.copy_src('x', 'm2')], [fn.sum('m', 'h_new'), fn.sum('m2', 'h_new2')], etype='follows')

#>>Not supported. Throws error: `TypeError: 'list' object is not callable`

[VoVAllen]

Hi, currently we won't support this fusion since we don't have much to improve on the performance.

[jermainewang]

Hi @fmellomascarenhas , to add more backgrounds to this, this was a feature before DGL v0.4. The motivation was, exactly as you mentioned, to open up possibility of fusing multiple message passing kernels together. We eventually removed such feature due to two reasons. First, it complicates the underlying implementation and most users tend to invoke multiple update_all for more readability. Second, such kind of kernel fusion can be eventually implemented in a much lower level such as on a dataflow graph or an IR produced by compilers like torchscript. This is a learned lesson for us too -- keeping every API as simple as possible and never over-designing.

[felipemello1]

Second, such kind of kernel fusion can be eventually implemented in a much lower level such as on a dataflow graph or an IR produced by compilers like torchscript. This is a learned lesson for us too -- keeping every API as simple as possible and never over-designing.

Hi @jermainewang , thank for explaining it. Would you have an example of this low level implementation? Is it using torch_scatter.scatter and working with the indices? Thanks!

[github-actions]

This issue has been automatically marked as stale due to lack of activity. It will be closed if no further activity occurs. Thank you

[github-actions]

This issue is closed due to lack of activity. Feel free to reopen it if you still have questions.