EnzymeAD / EnzymeAD/Enzyme

MLIR dataflow activity: a dup pointer used only through an aggregate is treated as inactive

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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.