mars-project / mars-project/mars

[PROPOSAL] Universal optimizer and smart fusion

Open
#593 0 comments 3 reactions 0 assignees View on GitHub
mod: optimization proposal
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.