[RFC] Compiling user-defined message/reduce functions with `torch.fx`
- Dominant language
- Python
- Stars
- 14.3k
- Forks
- 3.1k
- PR merge metrics
- No merged PRs in 30d
Description
## Motivation
DGL modules written w/ UDFs (User-Defined Functions) suffers severe performance issue, though we already recommend users to write their modules w/ built-in functions. Built-in functions are somewhat not intuitive and there are many legacy codes written w/ DGL UDFs.
There are a bunch of papers on compiling message-passing UDFs to existing sparse kernels, e.g. [Seastar](https://dl.acm.org/doi/10.1145/3447786.3456247), and [Graphiler](https://gnnsys.github.io/papers/GNNSys21_paper_10.pdf), Seastar was based on MindSpore ecosystem, and Graphiler was written upon Torchscript to manipulate the computational graph. However, Torchscript was not designed to transform IRs, and Graphiler needs user to change the syntax of their UDFs slightly.
Recently, [`torch.fx`](https://pytorch.org/docs/stable/fx.html) has brought a lot of attention because of its capability in symbolic tracing and transforming nn modules written in torch. There exists some work on [quantization](https://pytorch.org/tutorials/prototype/fx_graph_mode_quant_guide.html) and [kernel fusion](https://pytorch.org/tutorials/intermediate/fx_conv_bn_fuser.html) w/ `torch.fx`. More examples are available [here](https://github.com/pytorch/examples/tree/151944ecaf9ba2c8288ee550143ae7ffdaa90a80/fx).
## Pitch
I suppose `torch.fx` is a great fit for us to implement graphiler in DGL, for the following reasons:
1. Easy to use: most modules are exposed in Python side and well documented (see [GraphModule](https://pytorch.org/docs/stable/fx.html#torch.fx.GraphModule) and [Transformer](https://pytorch.org/docs/stable/fx.html#torch.fx.Transformer)), and we can easily manipulate the IR in pure Python.
2. `torch.fx` supports symbolic transformation without knowing about the input, we just need to override a `torch.fx.Transformer` module to perform the mp-UDF compilation; another mode [Interpreter](https://pytorch.org/docs/stable/fx.html#torch.fx.Interpreter) requires input tensor, which we can ignore for now.
3. Highly configurable, to support a torch NN module w/ DGLGraph's, a feasible solution is to [customize tracer](https://pytorch.org/docs/stable/fx.html#customizing-tracing). We can also enable graph-aware tracing, to enable more possible optimizations (e.g. AOT kernel-fusion and graph-aware kernel compilation w/ TVM TensorIR).
I'm working on a prototype PR on how to customize a tracer that recognizes DGLGraph and compile a simple GCN module written in UDFs.
Discussions are welcomed @VoVAllen @BarclayII @Rhett-Ying @jermainewang @xiezhq-hermann .
Contributor guide
No contributing guide indexed for this repository
Research direction
Start by reading the linked torch.fx documentation for GraphModule, Transformer, Interpreter, and customizing tracing, then review the proposed prototype for a custom tracer recognizing DGLGraph. Done would mean compiling a simple GCN module written with message/reduce UDFs, but the issue names no files or tests and remains an RFC.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- machine-learning
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Stale
- Clarity
- Needs clarification
- Newbie friendliness
- 25/100