Skip to content

Commit

Permalink
Extensions (#467)
Browse files Browse the repository at this point in the history 
* Use extensions for weakdeps

* LSP added "using Base: get_extension", which is probably a bad idea

* Fix extension syntax

* Fixes

* module

* Fix zygote tests

* Copy SIMDDualNumbers contents into extension, dropping it
  • Loading branch information
@chriselrod
chriselrod authored Feb 5, 2023
1 parent 807675d commit 57c7ffa
Show file tree
Hide file tree
Showing 3 changed files with 189 additions and 6 deletions.
4 changes: 1 addition & 3 deletions Project.toml
View file
Original file line number Diff line number Diff line change
@@ -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"
Expand All @@ -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"
Expand All @@ -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"
Expand Down
189 changes: 187 additions & 2 deletions ext/ForwardDiffExt.jl
View file
Original file line number Diff line number Diff line change
@@ -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,
Expand All @@ -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)
Expand Down
2 changes: 1 addition & 1 deletion src/predicates.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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])
Expand Down

2 comments on commit 57c7ffa

@chriselrod
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@JuliaRegistrator register

@JuliaRegistrator
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Registration pull request created: JuliaRegistries/General/77032

After the above pull request is merged, it is recommended that a tag is created on this repository for the registered package version.

This will be done automatically if the Julia TagBot GitHub Action is installed, or can be done manually through the github interface, or via:

git tag -a v0.12.150 -m "<description of version>" 57c7ffaf56eb3552e5584d9703de9115b9ad65a0
git push origin v0.12.150

Please sign in to comment.