EnzymeAD / EnzymeAD/Enzyme-JAX

enzyme-hlo-opt: simplify the raiser's masked lane pick (select/or reduce) over a pad mask to a slice

Open
#3,139 0 comments 0 reactions 0 assignees View on GitHub
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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.