EnzymeAD / EnzymeAD/Enzyme-JAX

EnzymeHLOOpt slice-motion pattern interplay produces type-inconsistent IR

Open
#3,000 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

Running `--enzyme-hlo-opt` (default options) on the module in
https://gist.github.com/wsmoses/86839597dfe6f13a3bcfa12ec7197572 (a raised mfem tetrahedron diffusion PA apply kernel, shapes already refined static) produces IR that fails verification:

```
enzymexlamlir-opt input.mlir --verify-each=0 --enzyme-hlo-opt -o out.mlir
enzymexlamlir-opt out.mlir -o /dev/null
# error: size of operand dimension 0 (3) is not equal to 1 or size of result dimension 0 (4)
```

The invalid output around the failure:

```mlir
%20 = stablehlo.constant dense<0.0> : tensor<3xf64>
%252 = stablehlo.slice %137 [0:3] : (tensor<4xf64>) -> tensor<3xf64>
%253 = stablehlo.multiply %20, %252 : tensor<3xf64>
%254 = stablehlo.broadcast_in_dim %253, dims = [0] : (tensor<3xf64>) -> tensor<4x2xf64> // invalid
%255 = stablehlo.multiply %252, %65 : (tensor<3xf64>, tensor<4xf64>) -> tensor<4xf64> // invalid
```

The input contains 4-lane tensors where only 3 lanes are live (`select(iota<4> < 3, x, 0)`-style padding masks, d1d=3 data in max-4 registers). Some pattern narrows values to the live 3 lanes and replaces uses whose consumers keep the padded 4-shape.

Bisection so far:
- `passses` bitmask bisect: only clearing **bit 0** (`SliceTransposeBase`, `SliceReshapeTranspose`, `SliceBroadcast`, `SliceReduceWindow`) avoids the invalid IR; clearing any other single bit still reproduces.
- Within that group, removing **`SliceBroadcast`** from the pattern set avoids it.
- A dimension-consistency assertion added inside `SliceBroadcast`'s own rewrite never fires — its rewrites are individually well-typed, so this is an interplay where its output enables another (possibly always-on) pattern that replaces a value with a differently-shaped one.
- A small hand-built `slice(mul(broadcast(select(mask, s, 0) * a), broadcast(b)))` MWE does **not** reproduce; the while-loop context in the full module appears necessary.

Found while running the MFEM CUDA suite through the raising path (#2968): the invalid module fails the exec-time verify and killed the test until a retry-without-bit-0 fallback was added on the consumer side (EnzymeAD/Reactant.jl#3223, `1d89255cd`).

🤖 Generated with [Claude Code](https://claude.com/claude-code)

https://claude.ai/code/session_016zErYp7upmqr4NHfhod9UD

Contributor guide

No contributing guide indexed for this repository

Research direction

Start by reproducing the issue with the linked gist using `enzy...-opt` and inspect the `SliceBroadcast` pattern together with the other bit-0 patterns. Compare the generated IR before the final verifier run, especially in the while-loop context; done means the default `--enzyme-hlo-opt` output passes the subsequent verification command without the consumer-side fallback.

Written by the indexing model from the issue text.

Assessment

Domain
compilers
Issue type
Bug
Difficulty
5/5
Estimated time
Over a week
Activity status
Active
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.