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

[wsmoses]

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

// 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<activity enzyme_dup>, #enzyme<activity enzyme_dup>],
      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:

%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

%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.

[pengmai]

The issue comes from this check, where the analysis determines if it should treat a value as active data or a pointer to active data. We're currently hardcoding a check if the value's type is an llvm.pointer or MemRef, but we're not considering structs. The longer term solution is to defer this check to an interface (probably AutoDiffTypeInterface), like so:

if (type.isMutable())
  joinActivePointerState(...)
else
  joinActiveDataState(...)

However, structs and composite types are a special case because they can contain active data, pointers, or both. The existing isMutable method is insufficient to capture that we want to call both join methods here. We could perhaps modify it to return a enum of "yes", "no", and "maybe" instead of a boolean yes/no.