EnzymeAD / EnzymeAD/Enzyme-JAX

One Hot Array Optimization Patterns

Open
#991 1 comment 0 reactions 0 assignees View on GitHub
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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.