The proper way to do automatic sparse differentiation on GPU

[ErikQQY]

I was trying to use DI+SMC+SCT on a small example on CUDA, not sure I am doing this right, really appreciate some suggestions.

using DifferentiationInterface
using SparseConnectivityTracer, SparseMatrixColorings
using SparseArrays, StableRNGs
using ForwardDiff
using ADTypes: ADTypes
using CUDA
const DI = DifferentiationInterface

n = 1000
J = spdiagm(0 => ones(n))

function MyAutoSparse(backend)
    sparsity_detector = ADTypes.KnownJacobianSparsityDetector(J)
    coloring_algorithm = GreedyColoringAlgorithm(RandomOrder(StableRNG(0)))
    return AutoSparse(backend; sparsity_detector, coloring_algorithm)
end

diffmode = MyAutoSparse(AutoForwardDiff())

x, y = CuArray(ones(n)), CuArray(zeros(n))

diffcache = DI.prepare_jacobian(copyto!, y, diffmode, x)
J = DI.jacobian(copyto!, y, diffcache, diffmode, x)

ERROR: LoadError: Scalar indexing is disallowed.
Invocation of getindex resulted in scalar indexing of a GPU array.
This is typically caused by calling an iterating implementation of a method.
Such implementations *do not* execute on the GPU, but very slowly on the CPU,
and therefore should be avoided.

If you want to allow scalar iteration, use `allowscalar` or `@allowscalar`
to enable scalar iteration globally or for the operations in question.
Stacktrace:
  [1] error(s::String)
    @ Base ./error.jl:44
  [2] errorscalar(op::String)
    @ GPUArraysCore ~/.julia/packages/GPUArraysCore/aNaXo/src/GPUArraysCore.jl:151
  [3] _assertscalar(op::String, behavior::GPUArraysCore.ScalarIndexing)
    @ GPUArraysCore ~/.julia/packages/GPUArraysCore/aNaXo/src/GPUArraysCore.jl:124
  [4] assertscalar(op::String)
    @ GPUArraysCore ~/.julia/packages/GPUArraysCore/aNaXo/src/GPUArraysCore.jl:112
  [5] getindex
    @ ~/.julia/packages/GPUArrays/3a5jB/src/host/indexing.jl:50 [inlined]
  [6] decompress!(A::SparseMatrixCSC{…}, B::CuArray{…}, result::SparseMatrixColorings.ColumnColoringResult{…})
    @ SparseMatrixColorings ~/.julia/packages/SparseMatrixColorings/cNs01/src/decompression.jl:379
  [7] _sparse_jacobian_aux!(::Tuple{…}, ::SparseMatrixCSC{…}, ::DifferentiationInterfaceSparseMatrixColoringsExt.SMCPushforwardSparseJacobianPrep{…}, ::AutoSparse{…}, ::CuArray{…})
    @ DifferentiationInterfaceSparseMatrixColoringsExt ~/.julia/packages/DifferentiationInterface/afUhd/ext/DifferentiationInterfaceSparseMatrixColoringsExt/jacobian.jl:326
  [8] jacobian!
    @ ~/.julia/packages/DifferentiationInterface/afUhd/ext/DifferentiationInterfaceSparseMatrixColoringsExt/jacobian.jl:237 [inlined]
  [9] jacobian(::typeof(copyto!), ::CuArray{…}, ::DifferentiationInterfaceSparseMatrixColoringsExt.SMCPushforwardSparseJacobianPrep{…}, ::AutoSparse{…}, ::CuArray{…})
    @ DifferentiationInterfaceSparseMatrixColoringsExt ~/.julia/packages/DifferentiationInterface/afUhd/ext/DifferentiationInterfaceSparseMatrixColoringsExt/jacobian.jl:250
 [10] top-level scope
    @ ~/work/demo.jl:35
 [11] include(mapexpr::Function, mod::Module, _path::String)
    @ Base ./Base.jl:307
 [12] top-level scope
    @ REPL[4]:1
in expression starting at /home/featurize/work/demo.jl:35
Some type information was truncated. Use `show(err)` to see complete types.

There are some scalar indexing errors with SparseMatrixColorings.jl when using the greedy coloring algorithm. I wonder if there should be some documentatios for this feature?

[gdalle]

Thanks for opening an issue, I agree we should document it. Can you paste the stack trace?

[gdalle]

At first glance I think you have to convert J to the proper cuSPARSE type at creation.

[ErikQQY]

Stacktrace updated!

Just turn J as a spdiagm with CuArray should be fine., for example

J = spdiagm(0 => CuArray(ones(n)))

works.

[gdalle]

Yeah, the sparsity pattern is taken as template to construct the actual Jacobian, so you need it to reflect the array type you expect

[ErikQQY]

Thanks for the prompt response. I was trying to test on a more complex ASD example using BandedMatrix, and it seems that BandedMatrices.jl has bad GPU support. Although it can somehow work, it still doesn't appear to integrate well with sparsity pattern detection. For example:

using DifferentiationInterface
using SparseConnectivityTracer, SparseMatrixColorings
using SparseArrays, StableRNGs
using ForwardDiff: ForwardDiff
using ADTypes: ADTypes
using CUDA
using BandedMatrices

const DI = DifferentiationInterface
n = 1000
# J = spdiagm(0 => CuArray(ones(n)), 1 => CuArray(ones(n - 1))) # spdiagm working fine
# Construct J according to https://github.com/JuliaLinearAlgebra/BandedMatrices.jl/issues/238
# basically _BandedMatrix(CuArray(ones(lower+upper+1, n)), n, lower, upper)
_J = _BandedMatrix(CuArray(ones(2, n)), n, 0, 1)
J = BandedMatrix(_J)

println("J constructed without error")

function MyAutoSparse(backend)
    sparsity_detector = ADTypes.KnownJacobianSparsityDetector(J)
    coloring_algorithm = GreedyColoringAlgorithm(RandomOrder(StableRNG(0)))
    return AutoSparse(backend; sparsity_detector, coloring_algorithm)
end

forward = MyAutoSparse(AutoForwardDiff())

x, y = CuArray(ones(n)), CuArray(zeros(n))

function f!(y, x)
    copyto!(y, x)
    copyto!(y, 2, x, 2, n - 1)
end

diffcache = DI.prepare_jacobian(f!, y, forward, x)
J_forward = DI.jacobian(f!, y, diffcache, forward, x)

Stacktrace:

Note the BandedMatrix on GPU can be successfully constructed

J constructed without error
ERROR: LoadError: Scalar indexing is disallowed.
Invocation of getindex resulted in scalar indexing of a GPU array.
This is typically caused by calling an iterating implementation of a method.
Such implementations *do not* execute on the GPU, but very slowly on the CPU,
and therefore should be avoided.

If you want to allow scalar iteration, use `allowscalar` or `@allowscalar`
to enable scalar iteration globally or for the operations in question.
Stacktrace:
  [1] error(s::String)
    @ Base ./error.jl:44
  [2] errorscalar(op::String)
    @ GPUArraysCore ~/.julia/packages/GPUArraysCore/aNaXo/src/GPUArraysCore.jl:151
  [3] _assertscalar(op::String, behavior::GPUArraysCore.ScalarIndexing)
    @ GPUArraysCore ~/.julia/packages/GPUArraysCore/aNaXo/src/GPUArraysCore.jl:124
  [4] assertscalar(op::String)
    @ GPUArraysCore ~/.julia/packages/GPUArraysCore/aNaXo/src/GPUArraysCore.jl:112
  [5] getindex
    @ ~/.julia/packages/GPUArrays/3a5jB/src/host/indexing.jl:50 [inlined]
  [6] scalar_getindex
    @ ~/.julia/packages/GPUArrays/3a5jB/src/host/indexing.jl:36 [inlined]
  [7] _getindex
    @ ~/.julia/packages/GPUArrays/3a5jB/src/host/indexing.jl:19 [inlined]
  [8] getindex
    @ ~/.julia/packages/GPUArrays/3a5jB/src/host/indexing.jl:17 [inlined]
  [9] inbands_getindex
    @ ~/.julia/packages/BandedMatrices/J4c91/src/banded/BandedMatrix.jl:414 [inlined]
 [10] banded_getindex
    @ ~/.julia/packages/BandedMatrices/J4c91/src/banded/BandedMatrix.jl:432 [inlined]
 [11] getindex
    @ ~/.julia/packages/BandedMatrices/J4c91/src/banded/BandedMatrix.jl:442 [inlined]
 [12] _getindex
    @ ./abstractarray.jl:1388 [inlined]
 [13] getindex
    @ ./abstractarray.jl:1342 [inlined]
 [14] iterate
    @ ./abstractarray.jl:1235 [inlined]
 [15] iterate
    @ ./abstractarray.jl:1233 [inlined]
 [16] SparseMatrixCSC{Float64, Int64}(M::BandedMatrix{Float64, CuArray{Float64, 2, CUDA.DeviceMemory}, Base.OneTo{Int64}})
    @ SparseArrays ~/.julia/juliaup/julia-1.12.6+0.x64.linux.gnu/share/julia/stdlib/v1.12/SparseArrays/src/sparsematrix.jl:891
 [17] SparseMatrixCSC
    @ ~/.julia/juliaup/julia-1.12.6+0.x64.linux.gnu/share/julia/stdlib/v1.12/SparseArrays/src/sparsematrix.jl:884 [inlined]
 [18] #BipartiteGraph#7
    @ ~/.julia/packages/SparseMatrixColorings/cNs01/src/graph.jl:355 [inlined]
 [19] BipartiteGraph
    @ ~/.julia/packages/SparseMatrixColorings/cNs01/src/graph.jl:354 [inlined]
 [20] _coloring(speed_setting::SparseMatrixColorings.WithResult, A::BandedMatrix{…}, ::ColoringProblem{…}, algo::GreedyColoringAlgorithm{…}, decompression_eltype::Type, symmetric_pattern::Bool; forced_colors::Nothing)
    @ SparseMatrixColorings ~/.julia/packages/SparseMatrixColorings/cNs01/src/interface.jl:237
 [21] _coloring
    @ ~/.julia/packages/SparseMatrixColorings/cNs01/src/interface.jl:227 [inlined]
 [22] #coloring#28
    @ ~/.julia/packages/SparseMatrixColorings/cNs01/src/interface.jl:196 [inlined]
 [23] coloring
    @ ~/.julia/packages/SparseMatrixColorings/cNs01/src/interface.jl:189 [inlined]
 [24] _prepare_sparse_jacobian_aux(::Val{…}, ::DifferentiationInterface.PushforwardFast, ::CuArray{…}, ::Tuple{…}, ::AutoSparse{…}, ::CuArray{…})
    @ DifferentiationInterfaceSparseMatrixColoringsExt ~/.julia/packages/DifferentiationInterface/afUhd/ext/DifferentiationInterfaceSparseMatrixColoringsExt/jacobian.jl:78
 [25] prepare_jacobian_nokwarg
    @ ~/.julia/packages/DifferentiationInterface/afUhd/ext/DifferentiationInterfaceSparseMatrixColoringsExt/jacobian.jl:57 [inlined]
 [26] #prepare_jacobian#47
    @ ~/.julia/packages/DifferentiationInterface/afUhd/src/first_order/jacobian.jl:18 [inlined]
 [27] prepare_jacobian(::typeof(f!), ::CuArray{…}, ::AutoSparse{…}, ::CuArray{…})
    @ DifferentiationInterface ~/.julia/packages/DifferentiationInterface/afUhd/src/first_order/jacobian.jl:15
 [28] top-level scope
    @ ~/work/demo.jl:32
 [29] include(mapexpr::Function, mod::Module, _path::String)
    @ Base ./Base.jl:307
 [30] top-level scope
    @ REPL[60]:1
in expression starting at /home/featurize/work/demo.jl:32
Some type information was truncated. Use `show(err)` to see complete types.

Although from testing, I can tell there are some issues with printing methods in banded matrices, I am not sure if SMC also have compatiable issues with BandedMatrices.jl?

[gdalle]

I'll take a look tomorrow

[gdalle]

Okay so I thought about it some more and it is perfectly expected that decompressing a GPU-based BandedMatrix doesn't work: there is no specialized code path for them, and so we fall back to the generic AbstractMatrix decompression with scalar indexing.
This might be motivation enough for me to finally solve JuliaDiff/SparseMatrixColorings.jl#64 by reworking JuliaDiff/SparseMatrixColorings.jl#139 so that it works on GPU too!

[ErikQQY]

Glad this is on the TODO list! Sorry for the naive question here, but doesn't JuliaDiff/SparseMatrixColorings.jl#139 state that BandedMatrix and BlockBandedMatrix don't need optimal decompressing on GPU? And from testing with that branch, such statements seem wrong and the scalar indexing issue still exists.

[gdalle]

I said it doesn't need specific decompression because the default happens to run on CPU and because I'm not sure it is that much slower. However, the fact that it fails on GPU is a good reason to add custom methods

[ErikQQY]

Oh, sorry I mean JuliaDiff/SparseMatrixColorings.jl#139. OK I see, so there are still some work to do for the full support of decompression of structured matrix on GPU. Is there anything I can help to push this forward, such as adding some methods or test on CUDA etc?