FluxML/Zygote.jl

fastmath support

Closed

#90 opened on Mar 6, 2019

 (2 comments) (0 reactions) (0 assignees)Julia (220 forks)batch import
help wanted

Repository metrics

Stars
 (1,568 stars)
PR merge metrics
 (Avg merge 2d 16h) (39 merged PRs in 30d)

Description

julia> foo(x) = 2.0x
foo (generic function with 1 method)

julia> bar(x) = @fastmath 2.0x
bar (generic function with 1 method)

julia> Zygote.derivative(foo, 1e6)
2.0

julia> Zygote.derivative(bar, 1e6)
ERROR: Non-differentiable function IntrinsicFunction
Stacktrace:
 [1] error(::String) at ./error.jl:33
 [2] (::Zygote.Pullback{Tuple{Core.IntrinsicFunction,Float64,Float64},Tuple{Core.IntrinsicFunction}})(::Int8) at /home/chriselrod/.julia/packages/Zygote/wYgcz/src/compiler/interface2.jl:0
 [3] mul_fast at ./fastmath.jl:163 [inlined]
 [4] (::Zygote.Pullback{Tuple{typeof(Base.FastMath.mul_fast),Float64,Float64},Tuple{Zygote.Pullback{Tuple{Core.IntrinsicFunction,Float64,Float64},Tuple{Core.IntrinsicFunction}}}})(::Int8) at /home/chriselrod/.julia/packages/Zygote/wYgcz/src/compiler/interface2.jl:0
 [5] bar at ./REPL[55]:1 [inlined]
 [6] (::Zygote.Pullback{Tuple{typeof(bar),Float64},Tuple{getfield(Zygote, Symbol("##265#back#163")){getfield(Zygote, Symbol("#back#162")){:mul_fast,Zygote.Context,Module,typeof(Base.FastMath.mul_fast)}},Zygote.Pullback{Tuple{typeof(Base.FastMath.mul_fast),Float64,Float64},Tuple{Zygote.Pullback{Tuple{Core.IntrinsicFunction,Float64,Float64},Tuple{Core.IntrinsicFunction}}}}}})(::Int8) at /home/chriselrod/.julia/packages/Zygote/wYgcz/src/compiler/interface2.jl:0
 [7] (::getfield(Zygote, Symbol("##73#74")){Zygote.Pullback{Tuple{typeof(bar),Float64},Tuple{getfield(Zygote, Symbol("##265#back#163")){getfield(Zygote, Symbol("#back#162")){:mul_fast,Zygote.Context,Module,typeof(Base.FastMath.mul_fast)}},Zygote.Pullback{Tuple{typeof(Base.FastMath.mul_fast),Float64,Float64},Tuple{Zygote.Pullback{Tuple{Core.IntrinsicFunction,Float64,Float64},Tuple{Core.IntrinsicFunction}}}}}}})(::Int8) at /home/chriselrod/.julia/packages/Zygote/wYgcz/src/compiler/interface.jl:38
 [8] gradient(::Function, ::Float64) at /home/chriselrod/.julia/packages/Zygote/wYgcz/src/compiler/interface.jl:44
 [9] derivative(::typeof(bar), ::Float64) at /home/chriselrod/.julia/packages/Zygote/wYgcz/src/compiler/interface.jl:47
 [10] top-level scope at REPL[57]:1

@fastmath is convenient for enabling the autovectorizer and contractions (ie, fma instructions) on unrolled expressions. Would be nice to add mul_fast and friends.

Contributor guide