dmlc / dmlc/dgl

[RFC] Compiling user-defined message/reduce functions with `torch.fx`

Open
#3,627 5 comments 0 reactions 0 assignees View on GitHub
feature request
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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.