feat: use native Mooncake HVP in forward-over-reverse

[gdalle]

Fixes #990 (can't use Mooncake.hessian yet because it is restrictive type-wise)

[codecov]

Codecov Report

❌ Patch coverage is 0% with 40 lines in your changes missing coverage. Please review.
✅ Project coverage is 26.65%. Comparing base (e1e49f1) to head (cd01a46).
⚠️ Report is 1 commits behind head on main.

Files with missing lines	Patch %	Lines
...ifferentiationInterfaceMooncakeExt/second_order.jl	0.00%	39 Missing ⚠️
...e/ext/DifferentiationInterfaceMooncakeExt/utils.jl	0.00%	1 Missing ⚠️

❗ There is a different number of reports uploaded between BASE (e1e49f1) and HEAD (cd01a46). Click for more details.

HEAD has 14 uploads less than BASE

Flag	BASE (e1e49f1)	HEAD (cd01a46)
DI	11	2
DIT	9	4

Additional details and impacted files

@@             Coverage Diff             @@
##             main    #1042       +/-   ##
===========================================
- Coverage   47.66%   26.65%   -21.01%     
===========================================
  Files         136      132        -4     
  Lines        8044     8006       -38     
===========================================
- Hits         3834     2134     -1700     
- Misses       4210     5872     +1662     

Flag	Coverage Δ
DI	6.75% <0.00%> (-27.74%)	⬇️
DIT	78.06% <ø> (-3.66%)	⬇️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[rsenne]

Tried out with CuArrays and its broken but i think its an upstream issue

using CUDA
using Mooncake
using DifferentiationInterface
using LinearAlgebra

const DI = DifferentiationInterface

logp(x) = -0.5f0 * sum(abs2, x)
x_gpu = CUDA.CuArray(randn(Float32, 8))
v_gpu = CUDA.CuArray(ones(Float32, 8))

cache = DI.prepare_hvp(logp, AutoMooncake(), x_gpu, (v_gpu,)) # fails

ERROR: ValueAndPullbackReturnTypeError: Found a value of type Ptr{CUDACore.CUstream_st} in output, but output is not permitted
│ to be or contain a pointer. This is because the amount of memory to which it refers is unknown, therefore Mooncake.jl
│ is unable to allocate appropriate memory for its gradients.
└

Stacktrace:
  [1] throw_forward_ret_type_error(y::Ptr{CUDACore.CUstream_st})
    @ Mooncake C:\Users\senne\.julia\packages\Mooncake\PWvVI\src\interface.jl:73
  [2] __exclude_unsupported_output_internal!(y::Ptr{CUDACore.CUstream_st}, ::Set{UInt64})
    @ Mooncake C:\Users\senne\.julia\packages\Mooncake\PWvVI\src\interface.jl:627
  [3] __exclude_unsupported_output_internal!(y::CuStream, address_set::Set{UInt64})
    @ Mooncake C:\Users\senne\.julia\packages\Mooncake\PWvVI\src\interface.jl:403
  [4] __exclude_unsupported_output_internal!(y::CUDACore.Managed{CUDACore.DeviceMemory}, address_set::Set{UInt64})
    @ Mooncake C:\Users\senne\.julia\packages\Mooncake\PWvVI\src\interface.jl:403
  [5] __exclude_unsupported_output_internal!(y::GPUArrays.RefCounted{CUDACore.Managed{…}}, address_set::Set{UInt64})
    @ Mooncake C:\Users\senne\.julia\packages\Mooncake\PWvVI\src\interface.jl:403
  [6] __exclude_unsupported_output_internal!(y::GPUArrays.DataRef{CUDACore.Managed{…}}, address_set::Set{UInt64})
    @ Mooncake C:\Users\senne\.julia\packages\Mooncake\PWvVI\src\interface.jl:403
  [7] __exclude_unsupported_output_internal!(y::CuArray{Float32, 1, CUDACore.DeviceMemory}, address_set::Set{UInt64})
    @ Mooncake C:\Users\senne\.julia\packages\Mooncake\PWvVI\src\interface.jl:403
  [8] __exclude_unsupported_output(y::CuArray{Float32, 1, CUDACore.DeviceMemory})
    @ Mooncake C:\Users\senne\.julia\packages\Mooncake\PWvVI\src\interface.jl:370
  [9] __exclude_func_with_unsupported_output(fx::Tuple{…})
    @ Mooncake C:\Users\senne\.julia\packages\Mooncake\PWvVI\src\interface.jl:378
 [10] prepare_pullback_cache(::Function, ::Vararg{Any}; config::Mooncake.Config)
    @ Mooncake C:\Users\senne\.julia\packages\Mooncake\PWvVI\src\interface.jl:646
 [11] prepare_pullback_cache
    @ C:\Users\senne\.julia\packages\Mooncake\PWvVI\src\interface.jl:640 [inlined]
 [12] prepare_pullback_nokwarg(::Val{…}, ::typeof(DifferentiationInterface.shuffled_gradient), ::AutoMooncake{…}, ::CuArray{…}, ::Tuple{…}, ::DifferentiationInterface.FunctionContext{…}, ::Constant{…}, ::Constant{…})
    @ DifferentiationInterfaceMooncakeExt C:\Users\senne\Documents\GitHub\DifferentiationInterface.jl\DifferentiationInterface\ext\DifferentiationInterfaceMooncakeExt\onearg.jl:14
 [13] _prepare_hvp_aux(::Val{…}, ::DifferentiationInterface.ReverseOverReverse, ::DifferentiationInterface.DontPrepareInner, ::typeof(logp), ::AutoMooncake{…}, ::CuArray{…}, ::Tuple{…})
    @ DifferentiationInterface C:\Users\senne\Documents\GitHub\DifferentiationInterface.jl\DifferentiationInterface\src\second_order\hvp.jl:663
 [14] prepare_hvp_nokwarg
    @ C:\Users\senne\Documents\GitHub\DifferentiationInterface.jl\DifferentiationInterface\src\second_order\hvp.jl:125 [inlined]
 [15] #prepare_hvp#65
    @ C:\Users\senne\Documents\GitHub\DifferentiationInterface.jl\DifferentiationInterface\src\second_order\hvp.jl:16 [inlined]
 [16] prepare_hvp(::typeof(logp), ::AutoMooncake{Nothing}, ::CuArray{Float32, 1, CUDACore.DeviceMemory}, ::Tuple{CuArray{…}})
    @ DifferentiationInterface C:\Users\senne\Documents\GitHub\DifferentiationInterface.jl\DifferentiationInterface\src\second_order\hvp.jl:8
 [17] top-level scope
    @ c:\Users\senne\Documents\GitHub\DifferentiationInterface.jl\ryan_test\DI_Mooncake_test.jl:12

As it fails with a diff error. Same one on PMCMC,

using Mooncake
using CUDA
using Random

logp(x) = -0.5f0 * sum(abs2, x)
x_gpu = CUDA.CuArray(randn(Float32, 8))
v_gpu = CUDA.CuArray(ones(Float32, 8))

cache = prepare_hvp_cache(logp, x_gpu)
value_and_hvp!!(cache, logp, v_gpu, x_gpu)

ERROR: Mooncake.IntrinsicsWrappers.MissingIntrinsicWrapperException("Unable to translate the intrinsic Val{Core.Intrinsics.atomic_pointerref}() into a regular Julia function. Please see github.com/chalk-lab/Mooncake.jl/issues/208 for more discussion.")
Stacktrace:
  [1] translate(f::Val{Core.Intrinsics.atomic_pointerref})
    @ Mooncake.IntrinsicsWrappers /projectnb/ssmsvi/rsenne/.julia/packages/Mooncake/O5brV/src/rules/builtins.jl:128
  [2] lift_intrinsic(::Core.IntrinsicFunction, ::Core.SSAValue, ::QuoteNode, ::Vararg{Any})
    @ Mooncake /projectnb/ssmsvi/rsenne/.julia/packages/Mooncake/O5brV/src/interpreter/ir_normalisation.jl:309
  [3] intrinsic_to_function
    @ /projectnb/ssmsvi/rsenne/.julia/packages/Mooncake/O5brV/src/interpreter/ir_normalisation.jl:300 [inlined]
  [4] normalise!(ir::Compiler.IRCode, spnames::Vector{Symbol})
    @ Mooncake /projectnb/ssmsvi/rsenne/.julia/packages/Mooncake/O5brV/src/interpreter/ir_normalisation.jl:35

But reading 1248 upstream this wouldn't work anyway yet

[gdalle]

I wasn't clear: when you use DI for a Mooncake HVP, you have to explicitly build a SecondOrder backend using the forward and the reverse object. This should give you the same error as the native API, which is only made for "forward over reverse"

[rsenne]

Oh my bad, also on me, I wasn't looking closely. I'll redo the example to ensure I get the same error

[rsenne]

Hrm actually getting a different error:

using CUDA
using Mooncake
using DifferentiationInterface
using LinearAlgebra

const DI = DifferentiationInterface

logp(x) = -0.5f0 * sum(abs2, x)
x_gpu = CUDA.CuArray(randn(Float32, 8))
v_gpu = CUDA.CuArray(ones(Float32, 8))

mk_so = SecondOrder(AutoMooncakeForward(), AutoMooncake())

cache = DI.prepare_hvp(logp, mk_so, x_gpu, (v_gpu,))
Hv_gpu = similar(x_gpu)

hvp_returned = DI.hvp!(
    logp,
    (Hv_gpu,),   
    cache,       
    mk_so,       
    x_gpu,       
    (v_gpu,),    
)

ERROR: UndefVarError: `value_and_hvp` not defined in `DifferentiationInterface`
Suggestion: check for spelling errors or missing imports.
Stacktrace:
 [1] getproperty
   @ .\Base_compiler.jl:47 [inlined]
 [2] gradient_and_hvp!(::typeof(logp), ::Tuple{…}, ::DifferentiationInterfaceMooncakeExt.MooncakeHVPPrep{…}, ::SecondOrder{…}, ::CuArray{…}, ::Tuple{…})
   @ DifferentiationInterfaceMooncakeExt C:\Users\senne\Documents\GitHub\DifferentiationInterface.jl\DifferentiationInterface\ext\DifferentiationInterfaceMooncakeExt\second_order.jl:62
 [3] hvp!(::typeof(logp), ::Tuple{…}, ::DifferentiationInterfaceMooncakeExt.MooncakeHVPPrep{…}, ::SecondOrder{…}, ::CuArray{…}, ::Tuple{…})
   @ DifferentiationInterfaceMooncakeExt C:\Users\senne\Documents\GitHub\DifferentiationInterface.jl\DifferentiationInterface\ext\DifferentiationInterfaceMooncakeExt\second_order.jl:89
 [4] top-level scope
   @ c:\Users\senne\Documents\GitHub\DifferentiationInterface.jl\ryan_test\DI_Mooncake_test.jl:17
Some type information was truncated. Use `show(err)` to see complete types.

[gdalle]

Woops, silly me. It should be better now