mlir: Func call reverse diff

[Pangoraw]

prompted by EnzymeAD/Reactant.jl#171 (comment).

We can now run enzyme-batch -> enzyme-wrap without inlining.

// ./bazel-bin/enzymexlamlir-opt test/lit_tests/batchtests/autodiff.mlir --enzyme-batch --enzyme-wrap="infn=main retTys=enzyme_active argTys=enzyme_active mode=ReverseModeCombined" --canonicalize --remove-unnecessary-enzyme-ops --arith-raise --cse --canonicalize --arith-raise --enzyme-simplify-math --mlir-print-ir-before=enzyme-batch
// -----// IR Dump Before BatchPass (enzyme-batch) //----- //
module {
  func.func @f(%arg0: tensor<f32>) -> tensor<f32> {
    %0 = stablehlo.multiply %arg0, %arg0 : tensor<f32>
    return %0 : tensor<f32>
  }
  func.func @main(%arg0: tensor<10xf32>) -> tensor<10xf32> {
    %0 = enzyme.batch @f(%arg0) {batch_shape = array<i64: 10>} : (tensor<10xf32>) -> tensor<10xf32>
    return %0 : tensor<10xf32>
  }
}


module {
  func.func @f(%arg0: tensor<f32>) -> tensor<f32> {
    %0 = stablehlo.multiply %arg0, %arg0 : tensor<f32>
    return %0 : tensor<f32>
  }
  func.func private @batched_f(%arg0: tensor<10xf32>) -> tensor<10xf32> {
    %0 = stablehlo.multiply %arg0, %arg0 : tensor<10xf32>
    return %0 : tensor<10xf32>
  }
  func.func @main(%arg0: tensor<10xf32>, %arg1: tensor<10xf32>) -> tensor<10xf32> {
    %cst = arith.constant dense<0.000000e+00> : tensor<10xf32>
    %0 = call @batched_f(%arg0) : (tensor<10xf32>) -> tensor<10xf32>
    %1 = stablehlo.add %arg1, %cst : tensor<10xf32>
    %2 = call @diffebatched_f(%arg0, %1) : (tensor<10xf32>, tensor<10xf32>) -> tensor<10xf32>
    %3 = stablehlo.add %2, %cst : tensor<10xf32>
    return %3 : tensor<10xf32>
  }
  func.func private @diffebatched_f(%arg0: tensor<10xf32>, %arg1: tensor<10xf32>) -> tensor<10xf32> {
    %cst = arith.constant dense<0.000000e+00> : tensor<10xf32>
    %0 = stablehlo.add %arg1, %cst : tensor<10xf32>
    %1 = stablehlo.multiply %0, %arg0 : tensor<10xf32>
    %2 = stablehlo.add %1, %cst : tensor<10xf32>
    %3 = stablehlo.add %2, %1 : tensor<10xf32>
    return %3 : tensor<10xf32>
  }
}

[wsmoses]

We should probably default to always returning the primal, no?

[Pangoraw]

I don't really understand how the primal would be used here

[wsmoses]

Yeah we would only need for fused forward and reverse passes or handling of mutation, we can defer that for later here

[wsmoses]

Long term We probably need to do the same augmented forward pass and separate reverse pass that enzyme llvm does atm .

if we assume everything in the function is read only and only return active not duplicated results this is fine

[Pangoraw]

ah right, the caching at the primal call site is not enough if the arguments are mutable. I could not make an example input with memrefs (something about invertPointerM). So I added an error if any of the result or arg of the initial call is mutable.

[wsmoses]

LGTM but some misc fix ups and also the assumptions of the current implementation should throw errors if not met (we have a isReadOnly for example)

[wsmoses]

Can this be a function interface instead?

[wsmoses]

Ah I see this use, this is super fragile. For example it won’t work if you have a calls b calls a.

I think this is okay to merge as is but adding proper recursion support shouldn’t be bad. We have it already for forward enzymemlir which basically just requires having a cache in enzymelogic asking if we’ve already made this derivative function

[Pangoraw]

yeah any cycle in the call graph will not be caught here.

[wsmoses]

Yeah enzymelogic should do this for forwards (and also llvm) already

[Pangoraw]

I will look at implementing something like ForwardCachedFunctions for reverse in a follow-up