EnzymeAD / EnzymeAD/Reactant.jl
Use Julia's type-inference to automatically detect which control-flow statements need tracing
- Dominant language
- Julia
- Stars
- 370
- Forks
- 74
- Avg merge
- 18h 47m
- Merged PRs (30d)
- 30
Description
So @jumerckx and I had this nice discussion on Slack in which we came to the conclusion that we can use Julia's type-inference to select which control-flow statements (I'm thinking of `if` but `for` or `while` could apply).
Based on @jumerckx work on IRCode introspection, we can do sth similar to @makslevental [`mlir-python-extra`](https://github.com/makslevental/mlir-python-extras)'s capability of converting Python `if`s constructs from the AST to MLIR `scf.if` ops. Although this is nice for a potential Brutus implementation, it doesn't completely fit Reactant's/XLA array processing model. The reason is that in Reactant's model, some statements in the traced code are considered "constant" from the MLIR's point of view. Basically, all code that doesn't deal with runtime values of traced `TracedRArray`s can safely be considered constant.
Reactant has a multistage compilation procedure:
1. We retrieve the IRCode with our own abstract interpreter (derived from Enzyme.jl) for replacing some methods
2. Currently not fully exploited but the idea is that control-flow analysis would run here with the IRCode from point (1)
3. We compile the IRCode and run the function, which emits the MLIR code
4. We command XLA to compile the MLIR and wrap the flattening/unflattening + compiled function call + sync calls in a `Thunk`
The idea is that in step (2), we would use the type-inferred type of the predicate of the conditionals to check whether we need to emit a `stablehlo.if` op on that conditional; i.e. `stablehlo.if` only should be emitted if the predicate on Julia `if` is a `TracedRNumber{Bool}`, so Julia's type-inference could tell us that for free.
@wsmoses later stated that there's another approach using Polygeist, which actually would allow us to get more Julia code into MLIR. I'm writing this here because it might be beneficial in mid-term, but mainly because it's a nice side-effect of using Julia for this and we don't want to forget it in case we need to revisit.
Contributor guide
No contributing guide indexed for this repository
Assessment
This issue has not been assessed yet.