OptArgs with GPU array

[jariji]

If I understand, fromrawu is supposed to change the Vector into a GPU array here but it doesn't: the first argument of f2 is based on a Vector.

using AccessibleOptimization, DifferentiationInterface, Optimization, OptimizationOptimJL, Random, GPUArrays, LinearAlgebra, AMDGPU, Zygote
let n = 8
    data = ROCArray(rand(n))
    f2(θ::AbstractGPUArray, x) = θ⋅x
    θ₀ = ROCArray(rand(n))
    vars = OptArgs(@o _[∗])
    optfunc = OptimizationFunction(f2, Optimization.AutoFiniteDiff())
    ops = OptProblemSpec((@o optfunc(_, data)), Vector, (θ₀), vars,)
    @time "solve" soln = solve(ops, BFGS(); iterations=1)
end

MethodError: no method matching (::var"#f2#88")(::SubArray{Float64, 1, Vector{Float64}, Tuple{UnitRange{Int64}}, true}, ::ROCArray{Float64, 1, AMDGPU.Runtime.Mem.HIPBuffer})

Closest candidates are:
  (::var"#f2#88")(::AbstractGPUArray, ::Any)

Am I supposed to do OptArgs(@o ROCArray(_[∗])) or overload fromrawu or _convert or are we waiting on Base.into or something else?

[aplavin]

Hmm, I can understand why you would expect f2 to be passed a gpu array here, but such preservation requires some care from all libraries involved.

Currently, Accessors.setall(oldarray, ::Elements, newarray) assumes that the user is happy to get either of the two array types back – either the old or the new one:
https://github.com/JuliaObjects/Accessors.jl/blob/d8a93147bfab5a918bf0e95c24ad3852c2631a94/src/getsetall.jl#L73
In your case, that's clearly not what you expect: your newarray is Vector as you requested Vector to be the underlying optimization vector type, but oldarray is a ROCArray that you expect to get back.

Not that details of getall/setall behavior, in particular container types, are explicitly considered experimental in Accessors and are subject to change. Basically, to allow fixing this kind of issues in the future :) Feel free to propose setall() improvements there!

[aplavin]

Trying out monkey-patching to see what overloads (setall and maybe others) are needed for your usecase could be helpful.

Am I supposed to do OptArgs(@o ROCArray(_[∗]))

First, this won't work as-is, because @o f(_[∗]) means "f applied to each element".
I guess the intent is captured by @o ROCArray(_)[∗], but you still shouldn't use it here. Basically, OptArgs describe how to get from your object θ (passed to f2) to a vector understandable by Optimization. This direction shouldn't include ROCArray call IIUC.

[jariji]

This works with AutoZygote and without passing a utype works. I'm not sure what the optimization array type is if I don't pass one.

using AccessibleOptimization, DifferentiationInterface, Optimization, OptimizationOptimJL, Random, GPUArrays, LinearAlgebra, AMDGPU, Zygote
let n = 8
    data = ROCArray(rand(n))
    f2(θ::AbstractGPUArray, x) = θ⋅x
    θ₀ = ROCArray(rand(n))
    vars = OptArgs(@o _[∗])
    optfunc = OptimizationFunction(f2, Optimization.AutoZygote())
    ops = OptProblemSpec((@o optfunc(_, data)), θ₀, vars,)
    @time "solve" soln = solve(ops, BFGS())
end

I tried playing with similar etc in

Accessors.setall(obj::AbstractArray, ::Elements, vs::ROCArray) = 
    (@assert length(obj) == length(vs); (reshape(vs, size(obj))))

but there are so many moving parts in this MWE with the autodiff and everything I kinda lose track of what's going on.

[aplavin]

without passing a utype works. I'm not sure what the optimization array type is if I don't pass one.

It should be automatic if not passed – basically, "whatever you get from Accessors.getall" :)

I have basically no experience working with GPU arrays (nor with reverse autodiff), probably won't be of much help here. But if you find some fixes or self-contained examples we can add to tests (even if @test_broken for now) – please report them!

Also, note that the main usecase for AccessibleOptimization (and Accessors in general, btw) is having some reasonably-small number of parameters that you specify with optics. Arrays and tuples also work of course, with _[∗], but large arrays (>> than current 8 elements) or many arrays can be suboptimal when mixed with other parameters.
This is not fundamental to the Accessors design, just the current implementation. The effort was mainly aimed towards making "small" (up to tens of values) stuff very cheap or zero-cost, with much less focus on larger arrays.