GPU scalar indexing when `transpose(::CuArray)`

[mattsignorelli]

It's weird because x = CUDA.zeros(1,4) works fine but x = transpose(CUDA.zeros(4)) gives the scalar indexing problem.

MWE:

using KernelAbstractions, CUDA
import DifferentiationInterface as DI

@kernel function foo!(y, x)
    i = @index(Global)
    a = 2*i - 1
    b = 2*i
    offset = (i-1)*4
    y[a] = (offset+1)*x[a] + (offset+2)*x[b]
    y[b] = (offset+3)*x[a] + (offset+4)*x[b]
end

kernel! = foo!(CUDA.CUDABackend())
f!(y,x) = kernel!(y, x, ndrange=2)

# This works fine:
x = CUDA.zeros(1,4)
y = CUDA.rand(1,4)
prep = DI.prepare_jacobian(f!, y, DI.AutoForwardFromPrimitive(DI.AutoForwardDiff()), x);
DI.value_and_jacobian!(fun!, y, jac, prep, DI.AutoForwardFromPrimitive(DI.AutoForwardDiff()), x)
#= Output:
4×4 CuArray{Float32, 2, CUDA.DeviceMemory}:
 1.0  2.0  0.0  0.0
 3.0  4.0  0.0  0.0
 0.0  0.0  5.0  6.0
 0.0  0.0  7.0  8.0
=#

# This causes scalar indexing:
x = transpose(CUDA.zeros(4))
y = transpose(CUDA.rand(4))
prep = DI.prepare_jacobian(f!, y, DI.AutoForwardFromPrimitive(DI.AutoForwardDiff()), x);
DI.value_and_jacobian!(fun!, y, jac, prep, DI.AutoForwardFromPrimitive(DI.AutoForwardDiff()), x)

[mattsignorelli]

jacobians are the same of course. The guilty lines seem to be in basis.jl, specifically

DifferentiationInterface.jl/DifferentiationInterface/src/utils/basis.jl

Line 18 in 7e80fef

b[i] = oneunit(eltype(b))

and

DifferentiationInterface.jl/DifferentiationInterface/src/utils/basis.jl

Line 36 in 7e80fef

b[i] = oneunit(eltype(b))

[gdalle]

What's going on here is that this method should be hit but isn't

DifferentiationInterface.jl/DifferentiationInterface/ext/DifferentiationInterfaceGPUArraysCoreExt/DifferentiationInterfaceGPUArraysCoreExt.jl

Lines 6 to 11 in 7e80fef

	function DI.basis(a::AbstractGPUArray{T}, i) where {T}
	b = similar(a)
	fill!(b, zero(T))
	@allowscalar b[i] = oneunit(T)
	return b
	end

The reason is that Transpose{GPUArray} is not itself a GPUArray, which should be easy to fix by widening the basis creation mechanism to include transpose and adjoint, as well as other wrappers.

More broadly however, the current implementation of AutoForwardFromPrimitive probably executes on the CPU, or with bad parallelism. I need to spend some more time thinking about it, to solve other issues like JuliaManifolds/ManifoldDiff.jl#84 too

[mattsignorelli]

I see. I was using AutoForwardFromPrimitive because of the comments in https://juliadiff.org/DifferentiationInterface.jl/DifferentiationInterface/stable/api/#From-primitive . Would it be more performant to instead allow scalar indexing and just use AutoForwardDiff?

[gdalle]

AutoForwardFromPrimitive is meant for this purpose, so you're not wrong to use it. Unlike ForwardDiff, at least it doesn't error with GPU arrays. However, I need to revisit it because it actually doesn't preserve the GPU nature of the arrays, and I think it is just running on the CPU then copying back.
It's a long-running PR (#620) which I still haven't found a satisfactory solution for.

[gdalle]

See #971

[mattsignorelli]

Hmm this example above actually does not error with AutoForwardDiff if I allow scalar indexing

[gdalle]

Yeah if you allow scalar indexing everything works either way

[mattsignorelli]

I tested this and yes, the actual differentiation is running on the CPU, regardless of transpose or not

[mattsignorelli]

Sorry I mispoke.. this is happening on the GPU actually

[gdalle]

Yeah so what happens is that

• there is a list of basis vectors, with type Vector{CuVector}
• they are iterated over on the CPU
• each one gives rise to a function run f(x::CuVector) on the GPU
• the results are gathered in a CPU Vector
• right now it is not transferred back to the GPU but it will be once fix: make wrong-mode pushforward/pullback return the correct array type #974 is merged

The issue with running every basis vector in parallel is that we would need to duplicate the prep object, and we would have two layers of parallelism, so I'm not sure it's the best idea.

Note that #974 won't solve the problem of Transpose and other wrappers, but it's a first step

[mattsignorelli]

As for the Transpose problem, one way is to instead dispatch on the CPU/GPU backend using KernelAbstractions.get_backend, which does seem to return the correct backend for the transposed GPU arrays.

BTW, I must say I am really pleased with this package, and want to thank you for all the work you've put into it. Even though it seems there are still some things being worked out with the GPU AD, I am currently doing massive batched Newton searches on the GPU using AutoSparse with this AutoForwardFromPrimitive(AutoForwardDiff()) It basically just worked out of the box. E.g. I can find the fixed points of 100,000 particles (6 phase space coordinates each) in 100,000 different systems, in just a few seconds, thanks to the AD tools provided in this package and then plugging in the Jacobians into CUBLAS's batched solvers. So thank you! :)

Code is here if curious

[gdalle]

Thank you for the kind words, it really means a lot! Besides citing our papers, it's the best thing I can ask for as a researcher ;)

Two things:

• If you're using AutoSparse, you're bypassing ForwardDiff's default Jacobian anyway, so I don't think you need to wrap the backend inside AutoForwardFromPrimitive for things to work on GPU.
• We added an extension to SparseMatrixColorings so that sparse autodiff could work directly on CUDA sparse matrices, wanna give it a try?