EnzymeAD / EnzymeAD/Enzyme-JAX
One Hot Array Optimization Patterns
- Dominant language
- MLIR
- Stars
- 131
- Forks
- 53
- Avg merge
- 1d 10h
- Merged PRs (30d)
- 193
Description
- [x] `scatter` with multiply (https://github.com/EnzymeAD/Enzyme-JAX/pull/1033)
```mlir
%cst_1 = stablehlo.constant 1 <.....>
%cst_0 = stablehlo.constant 0 <.....>
%materialized_one_hot_array = stablehlo.scatter(%cst_0, %cst_1)
... <- setindex essentially
%res = stablehlo.multiply %materialized_one_hot_array %b
```
convert this into a gather of indices of `%b` followed by a `scatter` of those indices
Contributor guide
No contributing guide indexed for this repository
Research direction
Start by locating the compiler pattern that recognizes the shown stablehlo.scatter followed by stablehlo.multiply sequence. Compare the existing scatter, gather, and setindex handling, then verify that the one-hot materialization is replaced by the requested gather-of-indices and scatter form without changing the result.
Written by the indexing model from the issue text.
Assessment
- Domain
- compilers
- Issue type
- Feature
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 35/100