Flux.jl icon indicating copy to clipboard operation
Flux.jl copied to clipboard

cpu() type stability

Open RobertGregg opened this issue 3 years ago • 2 comments

This is coming from a discussion I started on discourse. It looks like there is a type stability issue when using the cpu() function where the culprit seems to be Flux.fmap(). I don't think this is having a major impact on the performance of my model, but I wanted to document the issue just so you were aware of it.

Example of the issue:

julia> using Flux #v0.12.9

julia> @code_warntype cpu(rand(10))
MethodInstance for Flux.cpu(::Vector{Float64})
  from cpu(x) in Flux at ~/.julia/packages/Flux/qAdFM/src/functor.jl:146
Arguments
  #self#::Core.Const(Flux.cpu)
  x::Vector{Float64}
Locals
  #136::Flux.var"#136#137"
Body::Any
1 ─      (#136 = %new(Flux.:(var"#136#137")))
│   %2 = #136::Core.Const(Flux.var"#136#137"())
│   %3 = Flux.fmap(%2, x)::Any
└──      return %3

RobertGregg avatar Feb 18 '22 23:02 RobertGregg

We could add an overload for [cg]pu that skips the fmap path for arrays, but until it's demonstrated to cause actual performance issues I'm not sure that'll be much of a priority. 99% of the time there should be no need to call those functions on a hot code path, and the the VAE model zoo example you linked should be updated to reflect that (PRs welcome).

ToucheSir avatar Feb 18 '22 23:02 ToucheSir

Sounds fine to me. If I end up refactoring my VAE to not use cpu() in the main loop, I'll do a pull request for the model zoo.

RobertGregg avatar Feb 19 '22 00:02 RobertGregg