# How to use ForwardDiff with CuArrays

**URL:** <https://discourse.julialang.org/t/how-to-use-forwarddiff-with-cuarrays/39542>\
**Category:** GPU\
**Tags:** question, performance, differentiation, forwarddiff\
**Created:** [May 15, 2020, 6:03pm UTC](https://discourse.julialang.org/t/how-to-use-forwarddiff-with-cuarrays/39542 "2020-05-15T18:03:18Z")\
**Posts on this page:** 12\
**Page:** 1

<div class="post-metadata">

**Author:** ![clinton](https://avatars.discourse-cdn.com/v4/letter/c/7ba0ec/32.png) [@clinton](https://discourse.julialang.org/u/clinton)\
**Post date:** [May 15, 2020, 6:03pm UTC](https://discourse.julialang.org/t/how-to-use-forwarddiff-with-cuarrays/39542/1 "2020-05-15T18:03:18Z")

</div>

I want to use CuArrays to accelerate a Jacobian calculation for a function \Re^m \mapsto \Re^n where n\>\>m, which seems like a pretty good use case for ForwardDiff. I saw some usage of ForwardDiff in the tests for CuArrays, but I couldn’t figure out how to map what I was seeing in the test file to my (admittedly novice) usage of CuArrays, and there was no documentation on the subject that I could locate.

Here is an extremely simple mwe. IF I allow scalar getindex, it works yet performance seems atrocious (most likely due to user error). Note that the gradient crashes if `CuArrays.allowscalar(false)`, though simple evaluation of the function does not.

```nohighlight
using Revise
using Flux, BenchmarkTools, CuArrays, CUDAnative, ForwardDiff, LinearAlgebra, Random
CuArrays.allowscalar(true)
Random.seed!(1111)

function tcudiff(N, ::Type{T} = Float32) where T<:Real
  A = rand(T, N, N)

  cuA = A |> gpu

  f(A) = sum(A .+ A*A .+ T(1))

  @info "test f cpu: $(f(A))"
  (N<5) && @info "test ∇cpu: $(ForwardDiff.gradient(f, A))"
  @btime ForwardDiff.gradient($f, $A)

  @info "test f gpu: $(f(cuA))"
  (N<5) && @info "test ∇gpu: $(ForwardDiff.gradient(f, cuA))"
  @btime ForwardDiff.gradient($f, $cuA)
end

tcudiff(60)

```

Output:

```julia
[ Info: test f cpu: 60024.586
  351.611 ms (3004 allocations: 107.45 MiB)
[ Info: test f gpu: 60024.586
┌ Warning: Performing scalar operations on GPU arrays: This is very slow, consider disallowing these operations with `allowscalar(false)`
└ @ GPUArrays C:\Users\Clinton\.julia\packages\GPUArrays\JucdK\src\host\indexing.jl:43
  1.058 s (233939 allocations: 10.05 MiB)

```

What am I doing wrong?

---

<div class="post-metadata">

**Author:** ![ChrisRackauckas](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/chrisrackauckas/32/77_2.png) [@ChrisRackauckas](https://discourse.julialang.org/u/ChrisRackauckas)\
**Post date:** [May 15, 2020, 6:40pm UTC](https://discourse.julialang.org/t/how-to-use-forwarddiff-with-cuarrays/39542/2 "2020-05-15T18:40:09Z")

</div>

ForwardDiff’s seeding isn’t GPU-safe. But you can use `SparseDiffTools.jacobian` for this (and then transpose)

---

<div class="post-metadata">

**Author:** ![clinton](https://avatars.discourse-cdn.com/v4/letter/c/7ba0ec/32.png) [@clinton](https://discourse.julialang.org/u/clinton)\
**Post date:** [May 15, 2020, 8:21pm UTC](https://discourse.julialang.org/t/how-to-use-forwarddiff-with-cuarrays/39542/3 "2020-05-15T20:21:24Z")

</div>

I didn’t know that. Does using the SparseDiffTools.jacobian function make using the GPU kosher?

---

<div class="post-metadata">

**Author:** ![ChrisRackauckas](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/chrisrackauckas/32/77_2.png) [@ChrisRackauckas](https://discourse.julialang.org/u/ChrisRackauckas)\
**Post date:** [May 15, 2020, 8:29pm UTC](https://discourse.julialang.org/t/how-to-use-forwarddiff-with-cuarrays/39542/4 "2020-05-15T20:29:26Z")

</div>

Yeah. It’s using ForwardDiff under the hood, but directly using the duals. The issue is with how the duals are seeded in `ForwardDiff.jacobian`, while `SparseDiffTools.forward_color_jacobian` vectorizes this because it naturally works with color vectors for the seeding, which then has a side effect of working on GPUs. The sparse part of it doesn’t work well on GPUs since CuArrays.jl isn’t able to lower one of the broadcast expressions we use (broadcasted `getindex`), but hey, at least dense GPUs works well there!

---

<div class="post-metadata">

**Author:** ![clinton](https://avatars.discourse-cdn.com/v4/letter/c/7ba0ec/32.png) [@clinton](https://discourse.julialang.org/u/clinton)\
**Post date:** [May 16, 2020, 1:00am UTC](https://discourse.julialang.org/t/how-to-use-forwarddiff-with-cuarrays/39542/5 "2020-05-16T01:00:20Z")

</div>

This is great- the cached CPU version is quite speedy already. Using CuArrays, even cached, seems to generate more overhead than I expected, so I will need to see how my problem scales once I stop working with the MWE.

One small issue- using a cache with CuArrays generates a segfault on the out-of-place version. Not a big deal since the in-place version seems to work fine, though I am happy to file an issue if that would be useful.

---

<div class="post-metadata">

**Author:** ![ChrisRackauckas](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/chrisrackauckas/32/77_2.png) [@ChrisRackauckas](https://discourse.julialang.org/u/ChrisRackauckas)\
**Post date:** [May 16, 2020, 2:44am UTC](https://discourse.julialang.org/t/how-to-use-forwarddiff-with-cuarrays/39542/6 "2020-05-16T02:44:45Z")

</div>

> [@clinton](#):
>
> One small issue- using a cache with CuArrays generates a segfault on the out-of-place version. Not a big deal since the in-place version seems to work fine, though I am happy to file an issue if that would be useful.

Yup, open an issue please.

---

<div class="post-metadata">

**Author:** ![clinton](https://avatars.discourse-cdn.com/v4/letter/c/7ba0ec/32.png) [@clinton](https://discourse.julialang.org/u/clinton)\
**Post date:** [May 16, 2020, 3:54am UTC](https://discourse.julialang.org/t/how-to-use-forwarddiff-with-cuarrays/39542/7 "2020-05-16T03:54:23Z")

</div>

Done- see [108](https://github.com/JuliaDiff/SparseDiffTools.jl/issues/108)

---

<div class="post-metadata">

**Author:** ![RS-Coop](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/rs-coop/32/32175_2.png) [@RS-Coop](https://discourse.julialang.org/u/RS-Coop)\
**Post date:** [January 24, 2022, 5:09pm UTC](https://discourse.julialang.org/t/how-to-use-forwarddiff-with-cuarrays/39542/8 "2022-01-24T17:09:41Z")

</div>

I am running into a similar problem using ForwardDiff and Zygote for Hessian vector products, so I am unable to use SparseDiffTools functionality. Is there a reference you could point to for how I need to change my Dual seeding in order to make this work?

I looked at `forward_color_jacobian`, but it was unclear what I need to change. Perhaps this is because the changes are quite complicated, or introducing `Zygote.pullback` creates some issues.

---

<div class="post-metadata">

**Author:** ![ChrisRackauckas](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/chrisrackauckas/32/77_2.png) [@ChrisRackauckas](https://discourse.julialang.org/u/ChrisRackauckas)\
**Post date:** [January 24, 2022, 5:36pm UTC](https://discourse.julialang.org/t/how-to-use-forwarddiff-with-cuarrays/39542/9 "2022-01-24T17:36:56Z")

</div>

Hessian-vector products does not need forward\_color\_jacobian. You can just direct seed. Can you show what you ran?

---

<div class="post-metadata">

**Author:** ![RS-Coop](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/rs-coop/32/32175_2.png) [@RS-Coop](https://discourse.julialang.org/u/RS-Coop)\
**Post date:** [January 25, 2022, 12:34am UTC](https://discourse.julialang.org/t/how-to-use-forwarddiff-with-cuarrays/39542/10 "2022-01-25T00:34:03Z")

</div>

I was using the following code for the hvp operator. I thought I had fixed it by changing `one.(val)` to `one(val)` in the second to last line, but now I am having trouble reproducing the issue. Ignore the `Nothing` tag, I know I shouldn’t be doing that.

```julia
#=
In-place hvp operator compatible with Krylov.jl
=#
mutable struct HvpOperator{F, T, I}
	f::F
	x::AbstractArray{T, 1}
	dualCache1::AbstractArray{Dual{Nothing, T, 1}}
	size::I
	nProd::I
end

function HvpOperator(f, x::AbstractVector)
	dualCache1 = Dual.(x, similar(x))
	return HvpOperator(f, x, dualCache1, size(x, 1), 0)
end

Base.eltype(op::HvpOperator{F, T, I}) where{F, T, I} = T
Base.size(op::HvpOperator) = (op.size, op.size)

function LinearAlgebra.mul!(result::AbstractVector, op::HvpOperator, v::AbstractVector)
	op.nProd += 1

	op.dualCache1 .= Dual.(op.x, v)
	val, back = pullback(op.f, op.dualCache1)

	result .= partials.(back(one.(val))[1], 1)
end

```

In the GPU context `x` should be a CUDA array. Like I said, I am having some trouble reproducing the issue, but maybe it is possible to tell if my seeding process should be GPU compatible?

---

<div class="post-metadata">

**Author:** ![ChrisRackauckas](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/chrisrackauckas/32/77_2.png) [@ChrisRackauckas](https://discourse.julialang.org/u/ChrisRackauckas)\
**Post date:** [January 25, 2022, 1:28am UTC](https://discourse.julialang.org/t/how-to-use-forwarddiff-with-cuarrays/39542/11 "2022-01-25T01:28:29Z")

</div>

Oh this has nothing to do with the original post. Make a different thread.

---

<div class="post-metadata">

**Author:** ![RS-Coop](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/rs-coop/32/32175_2.png) [@RS-Coop](https://discourse.julialang.org/u/RS-Coop)\
**Post date:** [January 26, 2022, 3:55am UTC](https://discourse.julialang.org/t/how-to-use-forwarddiff-with-cuarrays/39542/12 "2022-01-26T03:55:13Z")

</div>

Currently having some gpu issues preventing me from recreating the issue, but if it comes up again I will do so! Sorry for extending this thread.