EnzymeAD / EnzymeAD/Enzyme-JAX
enzyme-hlo-opt: simplify the raiser's masked lane pick (select/or reduce) over a pad mask to a slice
- Dominant language
- MLIR
- Stars
- 131
- Forks
- 53
- Avg merge
- 1d 10h
- Merged PRs (30d)
- 193
Description
The raiser's racy-store refinement (#2949, `AffineToStableHLORaising.cpp` "pick any admitted lane's value") raises a store whose value varies along a parallel axis the destination does not index, under a mask, as a two-operand `stablehlo.reduce` picking one admitted lane:
```
%3:2 = stablehlo.reduce(%values init: %cst), (%mask init: %false) across dimensions = [0]
reducer(%av, %bv) (%am, %bm) { select(%am, %av, %bv); or(%am, %bm) }
```
When the mask is a `tid == k` guard it constant-folds to a pad of an all-true row (`pad %c, %false, high = [255, 0]` for 256 lanes), so the pick is lane k's value: `slice %values [k:k+1, ...]` reshaped, and the reduced mask is `true`. At 4096 × 256 (the `@lane0_large`-style kernel with `%v = in[e*256 + t]`) enzyme-hlo-opt today leaves the full reduction over a padded `256x4096xi1` mask and a transposed copy of the input:
```
%1 = stablehlo.pad %c, %false, low = [0, 0], high = [255, 0] : (tensor<1x4096xi1>, tensor) -> tensor<256x4096xi1>
%2 = stablehlo.transpose %0, dims = [1, 0] : (tensor<4096x256xf64>) -> tensor<256x4096xf64>
%3:2 = stablehlo.reduce(%2 init: %cst), (%1 init: %false) across dimensions = [0] ...
```
Which of several live lanes such a pick returns is unspecified (it is the reduction order), so any lane guaranteed true may be sliced: a padded index when the padding value is true and every reduced dim is padded, or the operand's first index when the operand is all true and the kept dims never read false padding. #3138 has a draft of this pattern (`OneHotMaskedReduce`), matching the select/or body explicitly; parking it here until the reduce shape itself is settled.
Part of #2968.
Contributor guide
No contributing guide indexed for this repository
Research direction
Start in AffineToStableHLORaising.cpp and compare the current raiser with the OneHotMaskedReduce draft referenced in #3138; inspect how enzyme-hlo-opt handles the @lane0_large-style shape. Done means the matching select/or masked reduction is replaced by a guaranteed-true lane slice without the padded mask reduction or transposed copy, once the reduce shape is settled.
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
- Active
- Clarity
- Mostly clear
- Newbie friendliness
- 38/100