Error in differentiating inv(matrix) with LAPACK function getrf!

[junyixu]

When using Trixi.jl with Enzyme.jl's automatic differentiation, I encountered an error with LAPACK's getrf! function:

ERROR: No augmented forward pass found for dgetrf_64_

The error occurs in the following function from Trixi.jl (src/solvers/dgsem/basis_lobatto_legendre.jl:771): junyixu/TrixiEnzyme.jl#1 (comment)

I've distilled the problem into a minimal working example:

MWE

using Enzyme

# Helper function for Legendre polynomials
function legendre_polynomial_and_derivative(n, x)
    if n == 0
        return 1.0, 0.0
    elseif n == 1
        return x, 1.0
    end
    
    p_prev = 1.0    # P₀(x)
    p = x           # P₁(x)
    dp_prev = 0.0   # P₀'(x)
    dp = 1.0        # P₁'(x)
    
    for i in 2:n
        p_new = ((2i-1)*x*p - (i-1)*p_prev)/i
        dp_new = ((2i-1)*(p + x*dp) - (i-1)*dp_prev)/i
        p_prev, dp_prev = p, dp
        p, dp = p_new, dp_new
    end
    return p, dp
end

# Main function to test
function vandermonde_legendre(nodes, N::Integer)
    n_nodes = length(nodes)
    n_modes = N + 1
    vandermonde = zeros(n_nodes, n_modes)
    for i in 1:n_nodes
        for m in 1:n_modes
            vandermonde[i, m], _ = legendre_polynomial_and_derivative(m - 1, nodes[i])
        end
    end
    return inv(vandermonde)  # only return inverse for simplification
end

# Test function using vandermonde_legendre
function f(x)
    nodes = [-1.0, 0.0, 1.0]  # fixed nodes for testing
    V_inv = vandermonde_legendre(nodes, 2)
    return sum(V_inv)  # return scalar output
end

# Try automatic differentiation
x = 1.0
result = autodiff(Reverse, f, Active, Active(1.0))

Environment

Enzyme v0.13.24
Trixi v0.9.13

I noticed a similar discussion in #1820 about LAPACK functions. Would defining custom Enzyme rules be a solution here, or are there other recommended approaches for handling LAPACK functions in automatic differentiation?

[junyixu]

Output of the MWE:

julia> include("troubleshooting.jl")
f (generic function with 1 method)

julia> result = autodiff(Reverse, f, Active, Active(1.0))
ERROR:
No augmented forward pass found for dgetrf_64_
 at context:   call void @dgetrf_64_(i8* nonnull %8, i8* nonnull %11, i64 %42, i8* nonnull %14, i64 %46, i64 %bitcast_coercion42) #77 [ "jl_roots"({} addrspace(10)* null, {} addrspace(10)* %39, {} addrspace(10)* null, {} addrspace(10)* %2, {} addrspace(10)* null, {} addrspace(10)* null) ], !dbg !145

Stacktrace:
 [1] #getrf!#1
   @ ~/.local/stow/julia-1.10.2/share/julia/stdlib/v1.10/LinearAlgebra/src/lapack.jl:565


Stacktrace:
  [1] #getrf!#1
    @ ~/.local/stow/julia-1.10.2/share/julia/stdlib/v1.10/LinearAlgebra/src/lapack.jl:565
  [2] getrf!
    @ ~/.local/stow/julia-1.10.2/share/julia/stdlib/v1.10/LinearAlgebra/src/lapack.jl:557 [inlined]
  [3] #lu!#158
    @ ~/.local/stow/julia-1.10.2/share/julia/stdlib/v1.10/LinearAlgebra/src/lu.jl:82 [inlined]
  [4] lu!
    @ ~/.local/stow/julia-1.10.2/share/julia/stdlib/v1.10/LinearAlgebra/src/lu.jl:81 [inlined]
  [5] #lu#164
    @ ~/.local/stow/julia-1.10.2/share/julia/stdlib/v1.10/LinearAlgebra/src/lu.jl:300 [inlined]
  [6] lu (repeats 2 times)
    @ ~/.local/stow/julia-1.10.2/share/julia/stdlib/v1.10/LinearAlgebra/src/lu.jl:299 [inlined]
  [7] inv
    @ ~/.local/stow/julia-1.10.2/share/julia/stdlib/v1.10/LinearAlgebra/src/dense.jl:917
  [8] vandermonde_legendre
    @ ~/WorkSpace/Enzyme_Trixi_origin/troubleshooting.jl:36
  [9] f
    @ ~/WorkSpace/Enzyme_Trixi_origin/troubleshooting.jl:42 [inlined]
 [10] diffejulia_f_161wrap
    @ ~/WorkSpace/Enzyme_Trixi_origin/troubleshooting.jl:0
 [11] macro expansion
    @ ~/.julia/packages/Enzyme/ydGh2/src/compiler.jl:5218 [inlined]
 [12] enzyme_call
    @ ~/.julia/packages/Enzyme/ydGh2/src/compiler.jl:4764 [inlined]
 [13] CombinedAdjointThunk
    @ ~/.julia/packages/Enzyme/ydGh2/src/compiler.jl:4636 [inlined]
 [14] autodiff
    @ ~/.julia/packages/Enzyme/ydGh2/src/Enzyme.jl:503 [inlined]
 [15] autodiff(mode::ReverseMode{false, false, FFIABI, false, false}, f::typeof(f), ::Type{Active}, args::Active{Float64})
    @ Enzyme ~/.julia/packages/Enzyme/ydGh2/src/Enzyme.jl:524
 [16] top-level scope
    @ REPL[2]:1

[wsmoses]

@michel2323 would you be able to look into adding an EnzymeRule for matrix inv?