From e06e0b636a4cc23b7011d1b0f0dd8079f39ec2d5 Mon Sep 17 00:00:00 2001 From: dan Date: Tue, 30 Jun 2026 06:39:02 +0200 Subject: [PATCH 1/2] enzyme: reverse-mode autodiff extension --- Project.toml | 3 + ext/EnzymeCoreExt.jl | 406 +++++++++++++++++++++++++++++++++++++++++++ ext/meta_kernels.jl | 41 +++++ test/Project.toml | 1 + test/enzyme.jl | 50 ++++++ test/runtests.jl | 11 ++ 6 files changed, 512 insertions(+) create mode 100644 ext/EnzymeCoreExt.jl create mode 100644 ext/meta_kernels.jl create mode 100644 test/enzyme.jl diff --git a/Project.toml b/Project.toml index 74c936ffa..29268afe9 100644 --- a/Project.toml +++ b/Project.toml @@ -36,9 +36,11 @@ Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2" UUIDs = "cf7118a7-6976-5b1a-9a39-7adc72f591a4" [weakdeps] +EnzymeCore = "f151be2c-9106-41f4-ab19-57ee4f262869" SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b" [extensions] +EnzymeCoreExt = "EnzymeCore" SpecialFunctionsExt = "SpecialFunctions" [compat] @@ -48,6 +50,7 @@ BFloat16s = "0.5, 0.6" CEnum = "0.4, 0.5" CodecBzip2 = "0.8.5" Crayons = "4" +EnzymeCore = "0.8.2" ExprTools = "0.1" GPUArrays = "11.5.6" GPUCompiler = "1.22.5" diff --git a/ext/EnzymeCoreExt.jl b/ext/EnzymeCoreExt.jl new file mode 100644 index 000000000..f70e09f87 --- /dev/null +++ b/ext/EnzymeCoreExt.jl @@ -0,0 +1,406 @@ +module EnzymeCoreExt + +using Metal +import Metal: mtlconvert, mtlfunction + +using EnzymeCore +using EnzymeCore.EnzymeRules +using GPUArrays + +include("meta_kernels.jl") + +# Inactive: these device-state accessors are data-independent (ObjC queue, task-local storage). +function EnzymeCore.EnzymeRules.inactive_noinl(::typeof(Metal.synchronize), args...) + return nothing +end +function EnzymeCore.EnzymeRules.inactive_noinl(::typeof(Metal.device), args...) + return nothing +end +function EnzymeCore.EnzymeRules.inactive_noinl(::typeof(Metal.global_queue), args...) + return nothing +end + +## mtlfunction + +function EnzymeCore.EnzymeRules.forward(config, ofn::Const{typeof(mtlfunction)}, + ::Type{<:Duplicated}, f::Const{F}, tt::Const{TT}; kwargs...) where {F,TT} + res = ofn.val(f.val, tt.val; kwargs...) + return Duplicated(res, res) +end + +function EnzymeCore.EnzymeRules.forward(config, ofn::Const{typeof(mtlfunction)}, + ::Type{BatchDuplicated{T,N}}, f::Const{F}, tt::Const{TT}; kwargs...) where {F,TT,T,N} + res = ofn.val(f.val, tt.val; kwargs...) + return BatchDuplicated(res, ntuple(_ -> res, Val(N))) +end + +function EnzymeCore.EnzymeRules.augmented_primal(config, ofn::Const{typeof(mtlfunction)}, + ::Type{RT}, f::Const{F}, tt::Const{TT}; kwargs...) where {F, CT, RT<:EnzymeCore.Annotation{CT}, TT} + res = ofn.val(f.val, tt.val; kwargs...) + primal = EnzymeRules.needs_primal(config) ? res : nothing + shadow = if EnzymeRules.needs_shadow(config) + EnzymeRules.width(config) == 1 ? res : ntuple(_ -> res, Val(EnzymeRules.width(config))) + else + nothing + end + return EnzymeRules.AugmentedReturn{EnzymeRules.primal_type(config, RT), EnzymeRules.shadow_type(config, RT), Nothing}(primal, shadow, nothing) +end + +function EnzymeCore.EnzymeRules.reverse(config, ofn::Const{typeof(mtlfunction)}, ::Type{RT}, subtape, f, tt; kwargs...) where {RT} + return (nothing, nothing) +end + +## mtlconvert + +function EnzymeCore.EnzymeRules.forward(config, ofn::Const{typeof(mtlconvert)}, ::Type{RT}, x::IT) where {RT, IT} + if EnzymeRules.needs_primal(config) && EnzymeRules.needs_shadow(config) + if EnzymeRules.width(config) == 1 + Duplicated(ofn.val(x.val), ofn.val(x.dval)) + else + BatchDuplicated(ofn.val(x.val), ntuple(i -> ofn.val(x.dval[i])::eltype(RT), Val(EnzymeRules.width(config)))) + end + elseif EnzymeRules.needs_shadow(config) + if EnzymeRules.width(config) == 1 + ofn.val(x.dval)::EnzymeCore.shadow_type(config, RT) + else + (ntuple(i -> ofn.val(x.dval[i])::eltype(RT), Val(EnzymeRules.width(config))))::EnzymeCore.shadow_type(config, RT) + end + elseif EnzymeRules.needs_primal(config) + ofn.val(x.val)::eltype(RT) + else + nothing + end +end + +function EnzymeCore.EnzymeRules.augmented_primal(config, ofn::Const{typeof(mtlconvert)}, ::Type{RT}, x::IT) where {RT, IT} + primal = EnzymeRules.needs_primal(config) ? ofn.val(x.val) : nothing + shadow = if EnzymeRules.needs_shadow(config) + EnzymeRules.width(config) == 1 ? ofn.val(x.dval) : + ntuple(i -> ofn.val(x.dval[i]), Val(EnzymeRules.width(config))) + else + nothing + end + return EnzymeRules.AugmentedReturn{EnzymeRules.primal_type(config, RT), EnzymeRules.shadow_type(config, RT), Nothing}(primal, shadow, nothing) +end + +function EnzymeCore.EnzymeRules.reverse(config, ofn::Const{typeof(mtlconvert)}, ::Type{RT}, tape, x::IT) where {RT, IT} + return (nothing,) +end + +## MtlArray constructors + +function EnzymeCore.EnzymeRules.forward(config, ofn::Const{Type{AT}}, + ::Type{RT}, uval::EnzymeCore.Annotation{UndefInitializer}, args...) where {AT <: MtlArray, RT} + primargs = ntuple(i -> args[i].val, Val(length(args))) + if EnzymeRules.needs_primal(config) && EnzymeRules.needs_shadow(config) + if EnzymeRules.width(config) == 1 + shadow = ofn.val(uval.val, primargs...)::AT + fill!(shadow, 0) + Duplicated(ofn.val(uval.val, primargs...), shadow) + else + tup = ntuple(Val(EnzymeRules.width(config))) do i + Base.@_inline_meta + shadow = ofn.val(uval.val, primargs...)::AT + fill!(shadow, 0) + shadow::AT + end + BatchDuplicated(ofn.val(uval.val, primargs...), tup) + end + elseif EnzymeRules.needs_shadow(config) + if EnzymeRules.width(config) == 1 + shadow = ofn.val(uval.val, primargs...)::AT + fill!(shadow, 0) + shadow::EnzymeCore.shadow_type(config, RT) + else + tup = ntuple(Val(EnzymeRules.width(config))) do i + Base.@_inline_meta + shadow = ofn.val(uval.val, primargs...)::AT + fill!(shadow, 0) + shadow::AT + end + tup::EnzymeCore.shadow_type(config, RT) + end + elseif EnzymeRules.needs_primal(config) + ofn.val(uval.val, primargs...) + else + nothing + end +end + +## make_zero / make_zero! + +@inline function EnzymeCore.make_zero( + x::MtlArray{FT}, +) where {FT<:AbstractFloat} + return Base.zero(x) +end +@inline function EnzymeCore.make_zero( + x::MtlArray{Complex{FT}}, +) where {FT<:AbstractFloat} + return Base.zero(x) +end + +@inline function EnzymeCore.make_zero( + ::Type{CT}, + seen::IdDict, + prev::CT, + ::Val{copy_if_inactive} = Val(false), +)::CT where {copy_if_inactive, FT<:AbstractFloat, CT <: Union{MtlArray{FT},MtlArray{Complex{FT}}}} + if haskey(seen, prev) + return seen[prev] + end + newa = Base.zero(prev) + seen[prev] = newa + return newa +end + +@inline function EnzymeCore.make_zero!( + prev::MtlArray{FT}, + seen::ST, +)::Nothing where {FT<:AbstractFloat,ST} + if !isnothing(seen) + if prev in seen + return nothing + end + push!(seen, prev) + end + fill!(prev, zero(FT)) + return nothing +end + +@inline function EnzymeCore.make_zero!( + prev::MtlArray{Complex{FT}}, + seen::ST, +)::Nothing where {FT<:AbstractFloat,ST} + if !isnothing(seen) + if prev in seen + return nothing + end + push!(seen, prev) + end + fill!(prev, zero(Complex{FT})) + return nothing +end + +## GPUArrays.mapreducedim! + +function EnzymeCore.EnzymeRules.forward(config, ofn::Const{typeof(GPUArrays.mapreducedim!)}, + ::Type{RT}, + f::EnzymeCore.Const{typeof(Base.identity)}, + op::EnzymeCore.Const{typeof(Base.add_sum)}, + R::EnzymeCore.Annotation{<:MtlArray{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 EnzymeRules.needs_primal(config) && EnzymeRules.needs_shadow(config) + R + elseif EnzymeRules.needs_shadow(config) + R.dval + elseif EnzymeRules.needs_primal(config) + R.val + else + nothing + 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{<:MtlArray{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{<:MtlArray{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 + +## GPUArrays._mapreduce + +function EnzymeCore.EnzymeRules.forward(config, ofn::Const{typeof(GPUArrays._mapreduce)}, + ::Type{RT}, + f::EnzymeCore.Const{typeof(Base.identity)}, + op::EnzymeCore.Const{typeof(Base.add_sum)}, + A::EnzymeCore.Annotation{<:MtlArray{T}}; dims::D, init) where {RT, T, D} + + if EnzymeRules.needs_primal(config) && EnzymeRules.needs_shadow(config) + if EnzymeRules.width(config) == 1 + shadow = ofn.val(f.val, op.val, A.dval; dims, init) + Duplicated(ofn.val(f.val, op.val, A.val; dims, init), shadow) + else + tup = ntuple(Val(EnzymeRules.batch_width(RT))) do i + Base.@_inline_meta + ofn.val(f.val, op.val, A.dval[i]; dims, init) + end + BatchDuplicated(ofn.val(f.val, op.val, A.val; dims, init), tup) + end + elseif EnzymeRules.needs_shadow(config) + if EnzymeRules.width(config) == 1 + ofn.val(f.val, op.val, A.dval; dims, init) + else + ntuple(Val(EnzymeRules.batch_width(RT))) do i + Base.@_inline_meta + ofn.val(f.val, op.val, A.dval[i]; dims, init) + end + end + elseif EnzymeRules.needs_primal(config) + ofn.val(f.val, op.val, A.val; dims, init) + else + nothing + 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{<:MtlArray{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{<:MtlArray{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 + +## HostKernel launch rules: differentiate the launch via device-side autodiff (meta_kernels.jl). + +# tape_type needs a GPU parent job; only its backend is used. +function EnzymeCore.compiler_job_from_backend(::Metal.MetalBackend, @nospecialize(F::Type), @nospecialize(TT::Type)) + mi = Metal.GPUCompiler.methodinstance(F, TT) + return Metal.GPUCompiler.CompilerJob(mi, Metal.compiler_config(Metal.device())) +end + +launch_extent(x::Integer) = Int(x) +launch_extent(x::NTuple{N, Integer}) where {N} = prod(Int, x) + +function EnzymeCore.EnzymeRules.forward(config, ofn::EnzymeCore.Annotation{Metal.HostKernel{F, TT}}, + ::Type{Const{Nothing}}, args...; + groups = 1, threads = 1, kwargs...) where {F, TT} + GC.@preserve args begin + dargs = ((mtlconvert(a) for a in args)...,) + TT2 = Tuple{typeof(config), F, map(typeof, dargs)...} + kernel = mtlfunction(metaf, TT2) + kernel(config, ofn.val.f, dargs...; groups, threads, kwargs...) + end + return nothing +end + +function EnzymeCore.EnzymeRules.augmented_primal(config, ofn::EnzymeCore.Annotation{Metal.HostKernel{F, TT}}, + ::Type{Const{Nothing}}, args0...; + groups = 1, threads = 1, kwargs...) where {F, TT} + args = ((mtlconvert(a) for a in args0)...,) + TapeType = EnzymeCore.tape_type( + EnzymeCore.compiler_job_from_backend(Metal.MetalBackend(), typeof(Base.identity), Tuple{Float32}), + ReverseSplitModified(EnzymeCore.set_runtime_activity(ReverseSplitWithPrimal, config), Val(EnzymeRules.overwritten(config))), + Const{F}, + Const{Nothing}, + map(typeof, args)..., + ) + n = launch_extent(threads) * launch_extent(groups) + subtape = Metal.MtlArray{TapeType}(undef, n) + + GC.@preserve args subtape begin + subtape2 = mtlconvert(subtape) + TT2 = Tuple{typeof(config), F, typeof(subtape2), map(typeof, args)...} + kernel = mtlfunction(meta_augf, TT2) + kernel(config, ofn.val.f, subtape2, args...; groups, threads, kwargs...) + end + + return EnzymeRules.AugmentedReturn{Nothing, Nothing, Metal.MtlArray}(nothing, nothing, subtape) +end + +function EnzymeCore.EnzymeRules.reverse(config, ofn::EnzymeCore.Annotation{Metal.HostKernel{F, TT}}, + ::Type{Const{Nothing}}, subtape, args0...; + groups = 1, threads = 1, kwargs...) where {F, TT} + args = ((mtlconvert(a) for a in args0)...,) + GC.@preserve args0 subtape begin + subtape2 = mtlconvert(subtape) + TT2 = Tuple{typeof(config), F, typeof(subtape2), map(typeof, args)...} + kernel = mtlfunction(meta_revf, TT2) + kernel(config, ofn.val.f, subtape2, args...; groups, threads, kwargs...) + end + return ntuple(Returns(nothing), Val(length(args0))) +end + +end # module EnzymeCoreExt diff --git a/ext/meta_kernels.jl b/ext/meta_kernels.jl new file mode 100644 index 000000000..741a75e03 --- /dev/null +++ b/ext/meta_kernels.jl @@ -0,0 +1,41 @@ +# Device-side meta kernels (mirrors AMDGPU.jl's ext/AMDGPUEnzymeCoreExt/meta_kernels.jl). + +@inline function thread_linear_index() + pos = Metal.thread_position_in_grid_3d() + tpg = Metal.threads_per_threadgroup_3d() + ng = Metal.threadgroups_per_grid_3d() + ex = tpg.x * ng.x + ey = tpg.y * ng.y + return (pos.x - 1) + ex * ((pos.y - 1) + ey * (pos.z - 1)) + 1 +end + +function metaf(config, fn, args::Vararg{Any, N}) where {N} + EnzymeCore.autodiff_deferred(EnzymeCore.set_runtime_activity(Forward, config), Const(fn), Const, args...) + return nothing +end + +function meta_augf(config, f, tape::Metal.MtlDeviceArray{TapeType}, args::Vararg{Any, N}) where {N, TapeType} + forward, _ = EnzymeCore.autodiff_deferred_thunk( + ReverseSplitModified(EnzymeCore.set_runtime_activity(ReverseSplitWithPrimal, config), Val(EnzymeRules.overwritten(config))), + TapeType, + Const{Core.Typeof(f)}, + Const{Nothing}, + map(typeof, args)..., + ) + i = thread_linear_index() + @inbounds tape[i] = forward(Const(f), args...)[1] + return nothing +end + +function meta_revf(config, f, tape::Metal.MtlDeviceArray{TapeType}, args::Vararg{Any, N}) where {N, TapeType} + _, reverse = EnzymeCore.autodiff_deferred_thunk( + ReverseSplitModified(EnzymeCore.set_runtime_activity(ReverseSplitWithPrimal, config), Val(EnzymeRules.overwritten(config))), + TapeType, + Const{Core.Typeof(f)}, + Const{Nothing}, + map(typeof, args)..., + ) + i = thread_linear_index() + @inbounds reverse(Const(f), args..., tape[i]) + return nothing +end diff --git a/test/Project.toml b/test/Project.toml index 46cfbb642..efa89cdfe 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -14,6 +14,7 @@ LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" Metal = "dde4c033-4e86-420c-a63e-0dd931031962" ObjectiveC = "e86c9b32-1129-44ac-8ea0-90d5bb39ded9" ParallelTestRunner = "d3525ed8-44d0-4b2c-a655-542cee43accc" +Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f" Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" ScopedValues = "7e506255-f358-4e82-b7e4-beb19740aa63" diff --git a/test/enzyme.jl b/test/enzyme.jl new file mode 100644 index 000000000..1b133becd --- /dev/null +++ b/test/enzyme.jl @@ -0,0 +1,50 @@ +using Test + +# Runs only when Enzyme is installed (opt-in via runtests.jl); mirrors AMDGPU's enzyme_tests.jl. +const enzyme_uuid = Base.UUID("7da242da-08ed-463a-9acd-ee780be4f1d9") +const enzyme_available = Base.locate_package(Base.PkgId(enzyme_uuid, "Enzyme")) !== nothing + +if enzyme_available + using EnzymeCore, Enzyme + import GPUCompiler + + function sq_kernel!(y, x) + i = Metal.thread_position_in_grid_1d() + @inbounds y[i] = x[i] * x[i] + return nothing + end + run_sq!(y, x) = (Metal.@metal threads = length(x) sq_kernel!(y, x); nothing) + + @testset "EnzymeCoreExt" begin + @test Base.get_extension(Metal, :EnzymeCoreExt) !== nothing + + @testset "compiler_job_from_backend" begin + @test EnzymeCore.compiler_job_from_backend( + Metal.MetalBackend(), typeof(() -> nothing), Tuple{}) isa GPUCompiler.CompilerJob + end + + @testset "sum gradient (vector)" begin + x = mtl(Float32[1, 2, 3, 4]) + dx = only(gradient(Reverse, sum, x)) + @test Array(dx) ≈ ones(Float32, 4) + end + + @testset "sum gradient (matrix)" begin + A = mtl(Float32[1 2 3; 4 5 6]) + dA = only(gradient(Reverse, sum, A)) + @test Array(dA) ≈ ones(Float32, 2, 3) + end + + @testset "kernel-launch gradient" begin + xv = Float32[1, 2, 3, 4] + x = mtl(xv) + y = mtl(zeros(Float32, 4)) + dx = mtl(zeros(Float32, 4)) + dy = mtl(ones(Float32, 4)) + Enzyme.autodiff(Reverse, run_sq!, Const, Duplicated(y, dy), Duplicated(x, dx)) + @test Array(dx) ≈ 2 .* xv # d/dx x^2 = 2x + end + end +else + @info "Skipping Enzyme tests (Enzyme not installed)" +end diff --git a/test/runtests.jl b/test/runtests.jl index 9b9ac5bb8..c8b2ae2ee 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -1,5 +1,6 @@ using Metal using ParallelTestRunner +import Pkg if !Metal.functional() @warn """Metal.jl is not functional on this system, so there is nothing to test; skipping. @@ -87,6 +88,16 @@ if filter_tests!(testsuite, args) if parse(Bool, get(ENV, "CI", "false")) || Sys.total_memory() < 12 * 2^30 delete!(testsuite, "largebroadcast") end + + # Enzyme tests are opt-in (Enzyme is a heavy, optional dependency); run via `runtests.jl enzyme`. + delete!(testsuite, "enzyme") +end + +# Add Enzyme / EnzymeCore on the fly when the tests are requested, rather than as test deps +# (matches AMDGPU.jl). +if any(name -> startswith(name, "enzyme"), keys(testsuite)) + @info "Running Enzyme tests" + Pkg.add(["EnzymeCore", "Enzyme"]) end # workers to run tests on From 458dd08948e453e20b188e38c472ddbe10e53dbf Mon Sep 17 00:00:00 2001 From: dan Date: Tue, 30 Jun 2026 06:39:44 +0200 Subject: [PATCH 2/2] enzyme: matrix-multiply rule --- ext/EnzymeCoreExt.jl | 19 +++++++++++++++++++ test/enzyme.jl | 11 +++++++++++ 2 files changed, 30 insertions(+) diff --git a/ext/EnzymeCoreExt.jl b/ext/EnzymeCoreExt.jl index f70e09f87..11b854e9a 100644 --- a/ext/EnzymeCoreExt.jl +++ b/ext/EnzymeCoreExt.jl @@ -6,6 +6,7 @@ import Metal: mtlconvert, mtlfunction using EnzymeCore using EnzymeCore.EnzymeRules using GPUArrays +using LinearAlgebra: mul! include("meta_kernels.jl") @@ -343,6 +344,24 @@ function EnzymeCore.EnzymeRules.reverse(config, ofn::Const{typeof(GPUArrays._map return (nothing, nothing, nothing) end +## Matrix multiply: MPS is opaque to Enzyme, so hook `*` with a closed-form reverse. + +function EnzymeCore.EnzymeRules.augmented_primal(config, ofn::Const{typeof(*)}, ::Type{RT}, + A::EnzymeCore.Annotation{<:MtlMatrix{T}}, B::EnzymeCore.Annotation{<:MtlMatrix{T}}) where {RT, T} + C = ofn.val(A.val, B.val) + dC = EnzymeRules.needs_shadow(config) ? fill!(similar(C), zero(T)) : nothing + primal = EnzymeRules.needs_primal(config) ? C : nothing + return EnzymeRules.AugmentedReturn(primal, dC, (A.val, B.val, dC)) +end + +function EnzymeCore.EnzymeRules.reverse(config, ofn::Const{typeof(*)}, ::Type{RT}, tape, + A::EnzymeCore.Annotation{<:MtlMatrix{T}}, B::EnzymeCore.Annotation{<:MtlMatrix{T}}) where {RT, T} + Aval, Bval, dC = tape + A isa Const || mul!(A.dval, dC, Bval', one(T), one(T)) + B isa Const || mul!(B.dval, Aval', dC, one(T), one(T)) + return (nothing, nothing) +end + ## HostKernel launch rules: differentiate the launch via device-side autodiff (meta_kernels.jl). # tape_type needs a GPU parent job; only its backend is used. diff --git a/test/enzyme.jl b/test/enzyme.jl index 1b133becd..2e340b29d 100644 --- a/test/enzyme.jl +++ b/test/enzyme.jl @@ -8,6 +8,8 @@ if enzyme_available using EnzymeCore, Enzyme import GPUCompiler + matmul_sum(A, B) = sum(A * B) + function sq_kernel!(y, x) i = Metal.thread_position_in_grid_1d() @inbounds y[i] = x[i] * x[i] @@ -35,6 +37,15 @@ if enzyme_available @test Array(dA) ≈ ones(Float32, 2, 3) end + @testset "matmul gradient" begin + A = mtl(Float32[1 2; 3 4]); B = mtl(Float32[5 6; 7 8]) + dA = mtl(zeros(Float32, 2, 2)); dB = mtl(zeros(Float32, 2, 2)) + Enzyme.autodiff(Enzyme.set_runtime_activity(Reverse), matmul_sum, Active, + Duplicated(A, dA), Duplicated(B, dB)) + @test Array(dA) ≈ Float32[11 15; 11 15] # d/dA sum(A*B) = ones*B' + @test Array(dB) ≈ Float32[4 4; 6 6] # d/dB sum(A*B) = A'*ones + end + @testset "kernel-launch gradient" begin xv = Float32[1, 2, 3, 4] x = mtl(xv)