DomCRose
commented
1 year ago xref: https://github.com/JuliaGPU/Adapt.jl/issues/21, https://github.com/JuliaGPU/CUDA.jl/issues/228
Since forward passes with nested wrappers will hit scalar indexing anyway, I think the best short term solution here would be to simply try and ensure that forward passes which only contain depth 1 wrappers, only result in depth 1 wrappers on the reverse pass for GPU arrays.
Materializing lazy array wrappers unecessarily could hamper CPU performance, so is it possible to add GPUArraysCore as a dependancy so that specialized methods can be added to ProjectTo to ensure this wrapper depth behaviour? Or would that increase load time too much?
Alternatively, perhaps the adjoint at https://github.com/JuliaDiff/ChainRulesCore.jl/blob/872d6452935abe754cece4fd96ca879f3a502b89/src/projection.jl#L413 could be replaced by a transpose and some sort of lazy conjugation, if such a thing is implemented anywhere?
As the title says. Code to see this:
Produces:
Interestingly, making a real and b complex allows it to run, but errors on display as the output type for the b gradient is becomes
Base.ReshapedArray{ComplexF64, 1, LinearAlgebra.Adjoint{ComplexF64, CuArray{ComplexF64, 2, CUDA.Mem.DeviceBuffer}}, Tuple{Base.MultiplicativeInverses.SignedMultiplicativeInverse{Int64}}}which it refuses to print. Collecting that array produces a CuArray with the correct gradient.The issue (at least with a complex and b real) seems to stem from https://github.com/JuliaDiff/ChainRulesCore.jl/blob/872d6452935abe754cece4fd96ca879f3a502b89/src/projection.jl#L413 creating a
Adjoint{ComplexF64, CuArray{ComplexF64, 2, CUDA.Mem.DeviceBuffer}}, which when then reshaped at https://github.com/JuliaDiff/ChainRulesCore.jl/blob/872d6452935abe754cece4fd96ca879f3a502b89/src/projection.jl#L230 creates aBase.ReshapedArray{ComplexF64, 1, Adjoint{ComplexF64, CuArray{ComplexF64, 2, CUDA.Mem.DeviceBuffer}}, Tuple{Base.MultiplicativeInverses.SignedMultiplicativeInverse{Int64}}}. Dispatch then sees this as an AbstractArray and sends it down base paths when map is called at https://github.com/JuliaDiff/ChainRulesCore.jl/blob/872d6452935abe754cece4fd96ca879f3a502b89/src/projection.jl#L236 rather than CUDA paths, resulting in scalar indexing.When a is real and b is complex, the element type of the gradient S matches the element type of the primal T, so the map in https://github.com/JuliaDiff/ChainRulesCore.jl/blob/872d6452935abe754cece4fd96ca879f3a502b89/src/projection.jl#L236 is not hit and instead the reshaped adjoint CuArray escapes, which I assume then hits scalar indexing show methods when dispatched for printing. However, in a more complicated function I guess this tangent would then enter later pullbacks and cause scalar indexing before gradient returns.
As far as I understand it, this would ideally be fixed by better wrapper array handling in Base / CuArray, but that seems like a hard and long lived issue. In the meantime I'm not sure what the best way to fix this would be, and whether that responsibility lies with CUDA or ChainRulesCore. Given the leaking of the wrapped array as a gradient of b in the a real, b complex case, perhaps there could be some tweaks to wrapped array handling here. Perhaps when the typeof dx is an Adjoint(...) then the reshape should be replaced by an adjoint followed by a broadcast of conj, or the earlier adjoint call in the ProjectTo{Adjoint} method should be a conj broadcast instead? Not sure what would be correct.