MLIR dataflow activity: a dup pointer used only through an aggregate is treated as inactive
- Dominant language
- LLVM
- Stars
- 1.7k
- Forks
- 188
- Avg merge
- 2d 4h
- Merged PRs (30d)
- 22
Description
On top of #3141 (forward mode for `llvm.extractvalue`/`llvm.insertvalue`), the same input differentiates correctly with the default activity analysis and incorrectly with `--enzyme=dataflow`: a `enzyme_dup` pointer whose uses flow through an aggregate is classified inactive, its shadow comes back null, and the generated tangent loads and stores through null pointers.
### MWE
```mlir
// enzymexlamlir-opt --enzyme %s -> correct
// enzymexlamlir-opt --enzyme=dataflow %s -> shadow is null
module {
func.func @square(%x : !llvm.ptr, %y : !llvm.ptr) {
%u = llvm.mlir.poison : !llvm.struct<(ptr, ptr)>
%a0 = llvm.insertvalue %x, %u[0] : !llvm.struct<(ptr, ptr)>
%a1 = llvm.insertvalue %y, %a0[1] : !llvm.struct<(ptr, ptr)>
%px = llvm.extractvalue %a1[0] : !llvm.struct<(ptr, ptr)>
%py = llvm.extractvalue %a1[1] : !llvm.struct<(ptr, ptr)>
%v = llvm.load %px : !llvm.ptr -> f64
%s = arith.mulf %v, %v : f64
llvm.store %s, %py : f64, !llvm.ptr
return
}
func.func @dsquare(%x : !llvm.ptr, %dx : !llvm.ptr, %y : !llvm.ptr, %dy : !llvm.ptr) {
enzyme.fwddiff @square(%x, %dx, %y, %dy) {
activity=[#enzyme, #enzyme],
ret_activity=[]
} : (!llvm.ptr, !llvm.ptr, !llvm.ptr, !llvm.ptr) -> ()
return
}
}
```
### Default activity: correct
The shadow aggregate carries the shadow pointers, the primal one the primal, and `2*x*dx` lands in `%dy`:
```mlir
%6 = llvm.insertvalue %arg1, %4[0] // shadow gets dx
%7 = llvm.insertvalue %arg0, %5[0] // primal gets x
%8 = llvm.insertvalue %arg3, %6[1] // shadow gets dy
%9 = llvm.insertvalue %arg2, %7[1] // primal gets y
...
llvm.store %18, %12 // tangent -> dy
llvm.store %19, %13 // primal -> y
```
### `dataflow`: shadow is null
```mlir
%1 = llvm.insertvalue %arg0, %0[0] // primal aggregate
%3 = llvm.mlir.zero : !llvm.ptr
%4 = llvm.insertvalue %3, %2[0] // shadow aggregate is all-null
%5 = llvm.mlir.zero : !llvm.ptr
%6 = llvm.insertvalue %5, %4[1]
%8 = llvm.extractvalue %6[0] // null
%12 = llvm.load %8 // load through null
...
llvm.store %16, %10 // store through null
```
`%arg1` and `%arg3` -- the shadows the caller passed for two `enzyme_dup` pointers -- are never used. `invertPointerM` on `%x` and `%y` returns a null shadow, so dataflow has concluded those arguments are inactive even though they are only ever read/written through the struct.
### Why it matters
This is the shape of every by-value capture: a lambda-taking kernel launcher (MFEM's `forall`, and the `CuWrap` templates under it) assembles the captured pointers into a struct and reads them back inside the kernel. With `-reactant-dataflow` on -- which is how the Reactant CUDA pipeline runs -- the MFEM dFEM GPU tests all come back with a zero tangent, which is what led here.
Contributor guide
Research direction
Start by running the provided MLIR reproducer with enzymexlamlir-opt using default activity and --enzyme=dataflow, then inspect the dataflow activity handling around invertPointerM and aggregate extractvalue/insertvalue operations. Done means enzyme_dup shadows remain available through the aggregate, the generated tangent reaches %dy, and no tangent load or store uses a null pointer.
Written by the indexing model from the issue text.
Assessment
- Domain
- compilers
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Active
- Clarity
- Mostly clear
- Newbie friendliness
- 45/100