mars-project / mars-project/mars
Use JAX to accelerate the execution of fuse operand
@wdkwyf is already working on this.
Since Aug 6, 2019.
- Dominant language
- Python
- Stars
- 2.7k
- Forks
- 325
- PR merge metrics
- No merged PRs in 30d
Description
**Is your feature request related to a problem? Please describe.**
We now use `numexpr` to speed up the execution of some operands, for example, `((a + 1) * 2).sum()` will be converted to expression string and called by `numexpr.evaluate()`.
`JAX` is another choice for optimization and cover more operands compared to `numexpr`, it can be used as one of the execution engine in `mars. optimizer`.
**Describe the solution you'd like**
The work can be split into three parts:
1. Implement `JaxFuseOperand` and `execute` method.
2. Add method `execute_jax` for each JAX supported operands. Not need to implement this on all the operands at first, we can choose arithmetic and reduction operands as an experiment.
3. Scan the subgraph and fuse all supported operands before executing fuse operand.
Some logic can refer to the implementation of `TensorNeFuseChunk`.
**Related issue**
#593
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Assessment
This issue has not been assessed yet.