Skip to content

Add Enzyme sum derivatives #2471

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 3 commits into from
Sep 4, 2024
Merged

Add Enzyme sum derivatives #2471

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
30d0420
Add Enzyme sum derivatives
wsmoses Aug 21, 2024
4a2199b
Update EnzymeCoreExt.jl
wsmoses Aug 22, 2024
84fd247
Update enzyme.jl
wsmoses Aug 22, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
159 changes: 159 additions & 0 deletions ext/EnzymeCoreExt.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@ else
using ..EnzymeCore
using ..EnzymeCore.EnzymeRules
end
using GPUArrays

function EnzymeCore.EnzymeRules.inactive(::typeof(CUDA.CUBLAS.handle))
return nothing
Expand Down Expand Up @@ -516,5 +517,163 @@ function EnzymeCore.EnzymeRules.noalias(::Type{CT}, ::UndefInitializer, args...)
return nothing
end

function EnzymeCore.EnzymeRules.forward(ofn::Const{typeof(GPUArrays.mapreducedim!)},
::Type{RT},
f::EnzymeCore.Const{typeof(Base.identity)},
op::EnzymeCore.Const{typeof(Base.add_sum)},
R::EnzymeCore.Annotation{<:AnyCuArray{T}}, A; init) where {RT, T}
if R isa Const || R isa Duplicated || R isa BatchDuplicated
ofn.val(f.val, op.val, R.val, A.val; init)
end

if A isa Duplicated || A isa DuplicatedNoNeed
if A isa Const
Base.fill!(R.dval, zero(T))
else
ofn.val(f.val, op.val, R.dval, A.dval)
end
elseif R isa BatchDuplicated || R isa BatchDuplicatedNoNeed
ntuple(Val(EnzymeRules.batch_width(R))) do i
Base.@_inline_meta
if A isa Const
Base.fill!(R.dval[i], zero(T))
else
ofn.val(f.val, op.val, R.dval[i], A.dval[i])
end
nothing
end
end

if RT <: Duplicated
return R
elseif RT <: Const
return R.val
elseif RT <: DuplicatedNoNeed
return R.dval
elseif RT <: BatchDuplicated
return R
else
return R.dval
end
end


function EnzymeCore.EnzymeRules.augmented_primal(config, ofn::Const{typeof(GPUArrays.mapreducedim!)},
::Type{RT},
f::EnzymeCore.Const{typeof(Base.identity)},
op::EnzymeCore.Const{typeof(Base.add_sum)},
R::EnzymeCore.Annotation{<:AnyCuArray{T}}, A; init) where {RT, T<:AbstractFloat}
if A isa Const || A isa Duplicated || A isa BatchDuplicated
ofn.val(f.val, op.val, R.val, A.val)
end

primal = if EnzymeRules.needs_primal(config)
R.val
else
nothing
end

shadow = if EnzymeRules.needs_shadow(config)
R.dval
else
nothing
end
return EnzymeRules.AugmentedReturn(primal, shadow, nothing)
end

function EnzymeCore.EnzymeRules.reverse(config, ofn::Const{typeof(GPUArrays.mapreducedim!)},
::Type{RT},
tape,
f::EnzymeCore.Const{typeof(Base.identity)},
op::EnzymeCore.Const{typeof(Base.add_sum)},
R::EnzymeCore.Annotation{<:AnyCuArray{T}}, A; init) where {RT, T<:AbstractFloat}

if !(A isa Const) && !(R isa Const)
if A isa Duplicated || A isa DuplicatedNoNeed
A.dval .+= R.dval
Base.fill!(R.dval, zero(T))
elseif A isa BatchDuplicated || A isa BatchDuplicatedNoNeed
ntuple(Val(EnzymeRules.batch_width(A))) do i
Base.@_inline_meta
A.dval[i] .+= R.dval[i]
Base.fill!(R.dval[i], zero(T))
nothing
end
end
end

return (nothing, nothing, nothing, nothing)
end

function EnzymeCore.EnzymeRules.forward(ofn::Const{typeof(GPUArrays._mapreduce)},
::Type{RT},
f::EnzymeCore.Const{typeof(Base.identity)},
op::EnzymeCore.Const{typeof(Base.add_sum)},
A::EnzymeCore.Annotation{<:AnyCuArray{T}}; dims::D, init) where {RT, T, D}
if RT <: Const
ofn.val(f.val, op.val, A.val; dims, init)
elseif RT <: Duplicated
(
ofn.val(f.val, op.val, A.val; dims, init),
ofn.val(f.val, op.val, A.dval; dims, init)
)
elseif RT <: DuplicatedNoNeed
ofn.val(f.val, op.val, A.dval; dims, init)
elseif RT <: BatchDuplicated
(
ofn.val(f.val, op.val, A.val; dims, init),
ntuple(Val(EnzymeRules.batch_width(RT))) do i
Base.@_inline_meta
ofn.val(f.val, op.val, A.dval[i]; dims, init)
end
)
else
@assert RT <: BatchDuplicatedNoNeed
ntuple(Val(EnzymeRules.batch_width(RT))) do i
Base.@_inline_meta
ofn.val(f.val, op.val, A.dval[i]; dims, init)
end
end
end

function EnzymeCore.EnzymeRules.augmented_primal(config, ofn::Const{typeof(GPUArrays._mapreduce)},
::Type{Active{RT}},
f::EnzymeCore.Const{typeof(Base.identity)},
op::EnzymeCore.Const{typeof(Base.add_sum)},
A::EnzymeCore.Annotation{<:AnyCuArray{T}}; dims::D, init) where {RT, T<:AbstractFloat, D}
primal = if EnzymeRules.needs_primal(config)
ofn.val(f.val, op.val, A.val; dims, init)
else
nothing
end

shadow = if EnzymeRules.needs_shadow(config)
A.dval
else
nothing
end
return EnzymeRules.AugmentedReturn(primal, shadow, nothing)
end

function EnzymeCore.EnzymeRules.reverse(config, ofn::Const{typeof(GPUArrays._mapreduce)},
dres::Active{RT},
tape,
f::EnzymeCore.Const{typeof(Base.identity)},
op::EnzymeCore.Const{typeof(Base.add_sum)},
A::EnzymeCore.Annotation{<:AnyCuArray{T}}; dims::D, init) where {RT, T<:AbstractFloat, D}

if A isa Duplicated || A isa DuplicatedNoNeed
A.dval .+= dres.val
elseif A isa BatchDuplicated || A isa BatchDuplicatedNoNeed
ntuple(Val(EnzymeRules.batch_width(A))) do i
Base.@_inline_meta
A.dval[i] .+= dres.val
nothing
end
end

return (nothing, nothing, nothing)
end

end # module

14 changes: 14 additions & 0 deletions test/extensions/enzyme.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -103,6 +103,20 @@ firstsum(x, y) = first(x .+ y)
#@test res[2] ≈ 1.2
end

@testset "Forward sum" begin
x = CuArray([1.0, 2.0, 3.0, 4.0])
dx = CuArray([100., 300.0, 500.0, 700.0])
res = Enzyme.autodiff(Forward, sum, Duplicated(x, dx))
@test res[1] ≈ 100+300+500+700.
end

@testset "Reverse sum" begin
x = CuArray([1.0, 2.0, 3.0, 4.0])
dx = CuArray([0., 0.0, 0.0, 0.0])
Enzyme.autodiff(Reverse, sum, Duplicated(x, dx))
@test all(dx .≈ 1.0)
end

# TODO once reverse kernels are in
# function togpu(x)
# x = CuArray(x)
Expand Down