mars-project / mars-project/mars
[PROPOSAL] Universal optimizer and smart fusion
- Dominant language
- Python
- Stars
- 2.7k
- Forks
- 325
- PR merge metrics
- No merged PRs in 30d
Description
# Background and Motivation
Mars has a lazy evaluation which means Mars will create a computational graph before execution, it allows optimization when we process and compile the graph.
Now Mars has a naive optimizer in the tensor module. The optimization can be divided into two parts:
1. Fusion on computational graph. Mars will fuse the nodes on a straight line into a single fuse node when build fine-grained graph.
2. Accelerating execution of fuse operand by `numexpr`.
As the boom of libraries like [JAX](https://github.com/google/jax) and [numba](https://github.com/numba/numba), we have more choices for the optimizer module, it's the time to make some enhancements on our optimizer.
There are two aspects to improve the performance, including graph fusion and optimizer for execution.
Fusing together adjacent nodes can provide several benefits:
1. Reduce the number of nodes in graph which lowers the scheduling overhead.
2. Improve the locality of data access, reduce the amount of data transmission.
3. Increases opportunities for optimization of execution, only fuse op could achieve performance acceleration by using `numexpr`, `JAX`, etc.
Meanwhile, more engines should be supported in `mars.Optimizer`, not only `numpexr` for now, as mentioned before, `JAX`, `numba` and `weld` can also be useful for our optimizer.
# Goals
- A smarter fusion strategy, not only a straight line fusion.
- Support more acceleration engine, optimized by them in a certain order.
- Both CPU and GPU support.
# Details
## Smarter Fusion
For now, straight line only fuse the Ops only have one inputs, such as unary Ops, reduction Ops and binary Ops with one input is a scalar, more rules are needed to fuse more nodes. listed below:
- Diamond
- Inverted Y-shaped
- ...
To fuse as many nodes as possible, we need scan the graph repeatedly, fuse straight line and
inverted Y-shape at first time, then fuse diamond in the graph, straight line again... after this cycle is repeated a few times, we can achieve a fully fused graph.
## Optimizer
Move `mars.Optimizer` out of tensor as it is helpful for both tensor and DataFrame modules. It can automatically choose a engine(numexpr, JAX...) for optimization when execute fuse node, or use all supported engines one by one in a certain order to maximizing optimization.
GPU support is also an important feature, our optimizer could detect data type and choose proper engines that support running on the GPU device. Furthermore, we should handle the situations that errors occurred when executed on GPU and degraded to CPU computation.
Contributor guide
Research direction
Start by locating the tensor module's existing optimizer and the proposed `mars.Optimizer` entry point, then map current straight-line fusion and numexpr execution. Done would require a defined strategy for repeated graph fusion, multiple execution engines, and CPU/GPU fallback, but the proposal names no files or tests.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- data, performance
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Stale
- Clarity
- Needs clarification
- Newbie friendliness
- 20/100