Extensions

Project.toml

@@ -1,7 +1,7 @@
 name = "LoopVectorization"
 uuid = "bdcacae8-1622-11e9-2a5c-532679323890"
 authors = ["Chris Elrod <[email protected]>"]
-version = "0.12.149"
+version = "0.12.150"
 
 [weakdeps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
@@ -28,7 +28,6 @@ LayoutPointers = "10f19ff3-798f-405d-979b-55457f8fc047"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 OffsetArrays = "6fe1bfb0-de20-5000-8ca7-80f57d26f881"
 PolyesterWeave = "1d0040c9-8b98-4ee7-8388-3f51789ca0ad"
-SIMDDualNumbers = "3cdde19b-5bb0-4aaf-8931-af3e248e098b"
 SIMDTypes = "94e857df-77ce-4151-89e5-788b33177be4"
 SLEEFPirates = "476501e8-09a2-5ece-8869-fb82de89a1fa"
 SnoopPrecompile = "66db9d55-30c0-4569-8b51-7e840670fc0c"
@@ -53,7 +52,6 @@ IfElse = "0.1"
 LayoutPointers = "0.1.11"
 OffsetArrays = "1.4.1"
 PolyesterWeave = "0.1.10, 0.2"
-SIMDDualNumbers = "0.1"
 SIMDTypes = "0.1"
 SLEEFPirates = "0.6.23"
 SnoopPrecompile = "1"

ext/ForwardDiffExt.jl

@@ -1,6 +1,10 @@
 module ForwardDiffExt
 import ForwardDiff, ChainRulesCore
-using SIMDDualNumbers, LoopVectorization
+using LoopVectorization, VectorizationBase, SLEEFPirates, ForwardDiff
+
+import IfElse: ifelse
+using VectorizationBase: AbstractSIMD, AbstractMask, zero_offsets
+
 using LoopVectorization:
   AbstractSIMD,
   AbstractStridedPointer,
@@ -18,7 +22,188 @@ using LoopVectorization:
   mask,
   vfnmadd_fast,
   mul_fast
-using VectorizationBase: zero_offsets
+
+@generated function Base.abs(
+  x::ForwardDiff.Dual{TAG,S,N}
+) where {TAG,S<:AbstractSIMD,N}
+  quote
+    $(Expr(:meta, :inline))
+    val = x.value
+    p = x.partials
+    cmp = val < zero($S)
+    absx = $ifelse(cmp, -val, val)
+    Base.Cartesian.@nexprs $N n -> p_n = p[n]
+    ForwardDiff.Dual{$TAG}(
+      absx,
+      ForwardDiff.Partials(
+        Base.Cartesian.@ntuple $N n -> $ifelse(cmp, -p_n, p_n)
+      )
+    )
+  end
+end
+@inline function Base.max(
+  x::ForwardDiff.Dual{TAG,<:AbstractSIMD,N},
+  y::ForwardDiff.Dual{TAG,<:AbstractSIMD,N}
+) where {TAG,N}
+  vx = ForwardDiff.value(x)
+  vy = ForwardDiff.value(y)
+  xgy = vx > vy
+  z = ifelse(xgy, vx, vy)
+  p = VectorizationBase.fmap(
+    ifelse,
+    xgy,
+    ForwardDiff.partials(x).values,
+    ForwardDiff.partials(y).values
+  )
+  ForwardDiff.Dual{TAG}(z, ForwardDiff.Partials(p))
+end
+
+@inline Base.max(
+  x::T,
+  y::Real
+) where {N,T<:ForwardDiff.Dual{<:Any,<:AbstractSIMD,N}} = max(x, T(y))
+@inline Base.max(
+  y::Real,
+  x::T
+) where {N,T<:ForwardDiff.Dual{<:Any,<:AbstractSIMD,N}} = max(x, T(y))
+@inline Base.max(
+  x::T,
+  y::Int
+) where {N,T<:ForwardDiff.Dual{<:Any,<:AbstractSIMD,N}} = max(x, T(y))
+@inline Base.max(
+  y::Int,
+  x::T
+) where {N,T<:ForwardDiff.Dual{<:Any,<:AbstractSIMD,N}} = max(x, T(y))
+
+@inline function Base.min(
+  x::ForwardDiff.Dual{TAG,<:AbstractSIMD,N},
+  y::ForwardDiff.Dual{TAG,<:AbstractSIMD,N}
+) where {TAG,N}
+  vx = ForwardDiff.value(x)
+  vy = ForwardDiff.value(y)
+  xgy = vx < vy
+  z = ifelse(xgy, vx, vy)
+  p = VectorizationBase.fmap(
+    ifelse,
+    xgy,
+    ForwardDiff.partials(x).values,
+    ForwardDiff.partials(y).values
+  )
+  ForwardDiff.Dual{TAG}(z, ForwardDiff.Partials(p))
+end
+@inline Base.min(
+  x::T,
+  y::Real
+) where {N,T<:ForwardDiff.Dual{<:Any,<:AbstractSIMD,N}} = min(x, T(y))
+@inline Base.min(
+  y::Real,
+  x::T
+) where {N,T<:ForwardDiff.Dual{<:Any,<:AbstractSIMD,N}} = min(x, T(y))
+@inline Base.min(
+  x::T,
+  y::Int
+) where {N,T<:ForwardDiff.Dual{<:Any,<:AbstractSIMD,N}} = min(x, T(y))
+@inline Base.min(
+  y::Int,
+  x::T
+) where {N,T<:ForwardDiff.Dual{<:Any,<:AbstractSIMD,N}} = min(x, T(y))
+
+@generated function SLEEFPirates.tanh_fast(
+  x::ForwardDiff.Dual{T,S,N}
+) where {T,S,N}
+  quote
+    $(Expr(:meta, :inline))
+    t = tanh_fast(x.value)
+    ∂t = $(VectorizationBase.vfnmadd_fast)(t, t, one(S))
+    p = x.partials
+    ForwardDiff.Dual{T}(
+      t,
+      ForwardDiff.Partials(
+        Base.Cartesian.@ntuple $N n -> $(Base.FastMath.mul_fast)(∂t, p[n])
+      )
+    )
+  end
+end
+@generated function SLEEFPirates.sigmoid_fast(
+  x::ForwardDiff.Dual{T,S,N}
+) where {T,S,N}
+  quote
+    $(Expr(:meta, :inline))
+    s = sigmoid_fast(x.value)
+    ∂s = $(VectorizationBase.vfnmadd_fast)(s, s, s)
+    p = x.partials
+    ForwardDiff.Dual{T}(
+      s,
+      ForwardDiff.Partials(
+        Base.Cartesian.@ntuple $N n -> $(Base.FastMath.mul_fast)(∂s, p[n])
+      )
+    )
+  end
+end
+@generated function VectorizationBase.relu(
+  x::ForwardDiff.Dual{T,S,N}
+) where {T,S,N}
+  quote
+    $(Expr(:meta, :inline))
+    v = x.value
+    z = zero(v)
+    cmp = v < z
+    r = ifelse(cmp, z, v)
+    p = x.partials
+    ForwardDiff.Dual{T}(
+      r,
+      ForwardDiff.Partials(Base.Cartesian.@ntuple $N n -> ifelse(cmp, z, p[n]))
+    )
+  end
+end
+
+@generated function ifelse(
+  m::AbstractMask,
+  x::ForwardDiff.Dual{TAG,V,P},
+  y::ForwardDiff.Dual{TAG,V,P}
+) where {TAG,V,P}
+  quote
+    $(Expr(:meta, :inline))
+    z = $ifelse(m, ForwardDiff.value(x), ForwardDiff.value(y))
+    px = ForwardDiff.partials(x)
+    py = ForwardDiff.partials(y)
+    p = Base.Cartesian.@ntuple $P p -> $ifelse(m, px[p], py[p])
+    ForwardDiff.Dual{$TAG}(z, ForwardDiff.Partials(p))
+  end
+end
+@generated function ifelse(
+  m::AbstractMask,
+  x::Number,
+  y::ForwardDiff.Dual{TAG,V,P}
+) where {TAG,V,P}
+  quote
+    $(Expr(:meta, :inline))
+    z = $ifelse(m, x, ForwardDiff.value(y))
+    py = ForwardDiff.partials(y)
+    p = Base.Cartesian.@ntuple $P p -> $ifelse(m, zero($V), py[p])
+    ForwardDiff.Dual{$TAG}(z, ForwardDiff.Partials(p))
+  end
+end
+@generated function ifelse(
+  m::AbstractMask,
+  x::ForwardDiff.Dual{TAG,V,P},
+  y::Number
+) where {TAG,V,P}
+  quote
+    $(Expr(:meta, :inline))
+    z = $ifelse(m, ForwardDiff.value(x), y)
+    px = ForwardDiff.partials(x)
+    p = Base.Cartesian.@ntuple $P p -> $ifelse(m, px[p], zero($V))
+    ForwardDiff.Dual{$TAG}(z, ForwardDiff.Partials(p))
+  end
+end
+@inline function SLEEFPirates.softplus(x::ForwardDiff.Dual{TAG}) where {TAG}
+  val = ForwardDiff.value(x)
+  expx = exp(val)
+  vx = log1p(expx)
+  px = Base.FastMath.inv_fast(one(val) + Base.FastMath.inv_fast(expx))
+  ForwardDiff.Dual{TAG}(vx, Base.FastMath.mul_fast(ForwardDiff.partials(x), px))
+end
 
 @generated function init_dual(v::Tuple{Vararg{AbstractSIMD,A}}) where {A}
   res = Expr(:tuple)

src/predicates.jl

@@ -11,7 +11,7 @@ isscopedname(:(Base.Checked.checked_add), (:Base, :Checked), :checked_add)
 function isscopedname(ex, modpath, name::Symbol)
   isexpr(ex, :(.), 2) &&
     (a = ex.args[2]; isa(a, QuoteNode) && a.value === name) &&
-    hasscope(ex.args[1], modpath)
+            hasscope(ex.args[1], modpath)
 end
 hasscope(modex, mod::Symbol) = modex === mod
 hasscope(modex, mod::Tuple{Symbol}) = hasscope(modex, mod[1])