EnzymeAD / EnzymeAD/Enzyme-JAX

raise-affine-to-stablehlo segfaults on affine.if yielding !llvm.ptr (pointer select)

Open
#2,928 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

With PR #2927 dissolving CUDA lambda-capture structs into scalars, MFEM's mass-integrator TU (\`fem/integ/bilininteg_mass_pa.cpp\`, xla-gpu backend) gets past all \`failed to raise\` errors and instead segfaults inside \`raise-affine-to-stablehlo\`.

The kernel region now contains a pointer select — MFEM's \`const_coeff ? C(0,0) : C(q,e)\` becomes an \`affine.if\` yielding a pointer (either the base or a GEP of it), which is then \`pointer2memref\`'d and loaded:

```mlir
%610 = affine.if #set1()[%193] -> !llvm.ptr {
affine.yield %324 : !llvm.ptr
} else {
%631 = llvm.getelementptr inbounds %324[%630] : (!llvm.ptr, i64) -> !llvm.ptr, !llvm.array<8 x i8>
affine.yield %631 : !llvm.ptr
}
%611 = "enzymexla.pointer2memref"(%610) : (!llvm.ptr) -> memref
%612 = affine.load %611[0] ...
```

Backtrace (release build, intermediate frames inlined):

```
#0 XXH3_hashLong_64b
#1 llvm::hash_value(llvm::StringRef)
#2 mlir::DenseStringElementsAttr::get(mlir::ShapedType, ArrayRef)
#3 mlir::DenseElementsAttr::get(mlir::ShapedType, ArrayRef)
#4 (anonymous namespace)::AffineToStableHLORaisingPass::runOnOperation()
```

\`DenseElementsAttr::get\` falling through to the string path means a splat/constant is being built whose shaped type has a non-int/float element type — consistent with an \`IntegerAttr::get(, ...)\` constructed somewhere along the pointer-typed if-result path (in a release build that produces garbage instead of asserting).

Reproduce: compile MFEM \`fem/integ/bilininteg_mass_pa.cpp\` with the Reactant plugin at \`-mllvm -reactant-backend -mllvm xla-gpu\` on top of #2927, or replay the imported module through the xla-gpu pipeline with \`enzymexlamlir-opt\`. A minimal \`cst ? c : c + i\` CUDA ternary does *not* reproduce it (that lowers to value select); the trigger is the structured \`affine.if -> !llvm.ptr\` with a GEP in one branch.

Expected: either raise the pointer select (select of the loaded values / offset select feeding the gather) or reject it with \`cannot raise op\` so \`err_if_not_fully_raised\` reports cleanly instead of crashing.

🤖 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 with the AffineToStableHLORaisingPass entry point and replay the imported module through the xla-gpu pipeline with enzymexlamlir-opt; the reproducing source is fem/integ/bilininteg_mass_pa.cpp. Focus on the affine.if yielding !llvm.ptr with a GEP branch, and consider the stated outcomes: raise the pointer select or reject it cleanly so err_if_not_fully_raised reports without a segfault.

Written by the indexing model from the issue text.

Assessment

Domain
backend, compilers
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Active
Clarity
Mostly clear
Newbie friendliness
55/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.