JuliaGPU / Adapt.jl

Other
86 stars 24 forks source link

`adapt_structure` can't handle `Box`es in closures #57

Closed ToucheSir closed 1 year ago

ToucheSir commented 1 year ago

Originally found in https://discourse.julialang.org/t/my-model-cant-be-transfered-to-gpu-adapt-jl-problem/92585/11. Copying over the MWE from that thread:

julia> function outer(a, c)
           b = 1
           inner(x) = x + a + b + c
           b = 2
           return inner
       end
outer (generic function with 1 method)

julia> clos = outer(1, 3)
(::var"#inner#1"{Int64, Int64}) (generic function with 1 method)

julia> dump(clos)
inner (function of type var"#inner#1"{Int64, Int64})
  a: Int64 1
  c: Int64 3
  b: Core.Box
    contents: Int64 2

julia> Adapt.adapt(Float64, clos)
ERROR: ArgumentError: tuple length should be ≥ 0, got -1
Stacktrace:
 [1] _ntuple(f::Adapt.var"#1#3"{var"#inner#1"{Int64, Int64}}, n::Int64)
   @ Base ./ntuple.jl:36
 [2] ntuple
   @ ./ntuple.jl:19 [inlined]
 [3] adapt_structure(to::Type, f::var"#inner#1"{Int64, Int64})
   @ Adapt ~/.julia/packages/Adapt/LAQOx/src/base.jl:17
 [4] adapt(to::Type, x::Function)
   @ Adapt ~/.julia/packages/Adapt/LAQOx/src/Adapt.jl:40
 [5] top-level scope
   @ REPL[7]:1

I imagine the logic in https://github.com/JuliaGPU/Adapt.jl/blob/v3.4.0/src/base.jl#L13-L21 needs to be adapted to handle nfields(f) > npar.

vchuravy commented 1 year ago

We won't be able to handle Core.Box it's semantics are incompatible with GPU execution.

ToucheSir commented 1 year ago

Certainly. For context, this is for pure host-side functions like Flux.{cpu,gpu,f32,f64}, which use Adapt under the hood.

maleadt commented 1 year ago

Please test https://github.com/JuliaGPU/Adapt.jl/pull/58

ToucheSir commented 1 year ago

Sorry, didn't get a chance to test #58 until today but it looks good! Also learned quite a bit about function type representation from reading the code :)