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