From 763dcec654fe3ddf1d755a575698da73677ad9cf Mon Sep 17 00:00:00 2001 From: Ryan Senne Date: Sat, 6 Jun 2026 17:56:05 -0400 Subject: [PATCH 1/3] Add linalg support via GPUArrays extension --- ext/EnzymeGPUArraysCoreExt.jl | 402 ++++++++++++++++++++++++++++++++++ test/cuda.jl | 87 +++++++- test/ext/jlarrays.jl | 77 +++++++ 3 files changed, 564 insertions(+), 2 deletions(-) diff --git a/ext/EnzymeGPUArraysCoreExt.jl b/ext/EnzymeGPUArraysCoreExt.jl index 12e3417338..b9548eb249 100644 --- a/ext/EnzymeGPUArraysCoreExt.jl +++ b/ext/EnzymeGPUArraysCoreExt.jl @@ -2,6 +2,18 @@ module EnzymeGPUArraysCoreExt using GPUArraysCore using Enzyme +using LinearAlgebra: LinearAlgebra, mul!, dot, Transpose, Adjoint, adjoint +using Enzyme.EnzymeCore: EnzymeCore +using Enzyme.EnzymeCore.EnzymeRules: + EnzymeRules, + FwdConfig, + RevConfig, + Annotation, + AugmentedReturn, + needs_primal, + needs_shadow, + overwritten, + width function Enzyme.zerosetfn(x::AbstractGPUArray, i::Int) res = zero(x) @@ -35,4 +47,394 @@ end end end +@inline _bget(x, ::Val{1}, ::Int) = x +@inline _bget(x, ::Val{N}, i::Int) where {N} = x[i] + +# Project an accumulated cotangent onto the (possibly real) parameter type. +_project(::Type{<:Real}, x) = real(x) +_project(::Type, x) = x + +# A GPU array or a (lazy) transpose/adjoint of one; the operand types that show +# up in `A * B` matmuls, including `transpose(X) * y`. +const MaybeWrappedGPU = Union{ + AbstractGPUArray, + Transpose{<:Any, <:AbstractGPUArray}, + Adjoint{<:Any, <:AbstractGPUArray}, +} + +#= +mul!(C, A, B, α, β): C = α·A·B + β·C₀ + +JVP: dC = α·(dA·B + A·dB) + β·dC₀ (+ dα·A·B + dβ·C₀) +Pullback: dA += conj(α)·dC·B' + dB += conj(α)·A'·dC + dC := conj(β)·dC (cotangent w.r.t. C₀) + dα = conj(⟨dC, A·B⟩) + dβ = conj(⟨dC, C₀⟩) +=# + +function EnzymeRules.forward( + config::FwdConfig, + func::Const{typeof(mul!)}, + RT::Type{<:Annotation}, + C::Annotation{<:AbstractGPUArray}, + A::Annotation, + B::Annotation, + α::Annotation{<:Number}, + β::Annotation{<:Number}, + ) + N = width(config) + if !(C isa Const) + # Update each output tangent from the old tangent values + # before the in-place primal overwrites C.val. + ntuple(Val(N)) do b + Base.@_inline_meta + dC = _bget(C.dval, Val(N), b) + tmp = β.val .* dC + if !(A isa Const) + tmp = tmp .+ α.val .* (_bget(A.dval, Val(N), b) * B.val) + end + if !(B isa Const) + tmp = tmp .+ α.val .* (A.val * _bget(B.dval, Val(N), b)) + end + if !(α isa Const) + tmp = tmp .+ _bget(α.dval, Val(N), b) .* (A.val * B.val) + end + if !(β isa Const) + tmp = tmp .+ _bget(β.dval, Val(N), b) .* C.val + end + copyto!(dC, tmp) + nothing + end + end + + func.val(C.val, A.val, B.val, α.val, β.val) + + if needs_primal(config) && needs_shadow(config) + return N == 1 ? Duplicated(C.val, C.dval) : BatchDuplicated(C.val, C.dval) + elseif needs_shadow(config) + return C.dval + elseif needs_primal(config) + return C.val + else + return nothing + end +end + +function EnzymeRules.augmented_primal( + config::RevConfig, + func::Const{typeof(mul!)}, + ::Type{RT}, + C::Annotation{<:AbstractGPUArray}, + A::Annotation, + B::Annotation, + α::Annotation{<:Number}, + β::Annotation{<:Number}, + ) where {RT} + # C₀ is needed for dβ; copy it before the primal overwrites C.val. + cache_C = !(β isa Const) ? copy(C.val) : nothing + + func.val(C.val, A.val, B.val, α.val, β.val) + + primal = needs_primal(config) ? C.val : nothing + shadow = needs_shadow(config) ? C.dval : nothing + + # A is needed for dB, B for dA; cache them only if overwritten before reverse. + cache_A = (overwritten(config)[3] && !(B isa Const) && !(C isa Const)) ? copy(A.val) : nothing + cache_B = (overwritten(config)[4] && !(A isa Const) && !(C isa Const)) ? copy(B.val) : nothing + cache_α = !(α isa Const) ? A.val * B.val : nothing + + return AugmentedReturn(primal, shadow, (cache_C, cache_A, cache_B, cache_α)) +end + +function EnzymeRules.reverse( + config::RevConfig, + func::Const{typeof(mul!)}, + ::Type{RT}, + tape, + C::Annotation{<:AbstractGPUArray}, + A::Annotation, + B::Annotation, + α::Annotation{<:Number}, + β::Annotation{<:Number}, + ) where {RT} + cache_C, cache_A, cache_B, cache_α = tape + Cval = cache_C !== nothing ? cache_C : C.val + Aval = cache_A !== nothing ? cache_A : A.val + Bval = cache_B !== nothing ? cache_B : B.val + N = width(config) + + if !(C isa Const) + dα = if !(α isa Const) + if N == 1 + _project(typeof(α.val), conj(dot(C.dval, cache_α))) + else + ntuple(i -> _project(typeof(α.val), conj(dot(C.dval[i], cache_α))), Val(N)) + end + else + nothing + end + dβ = if !(β isa Const) + if N == 1 + _project(typeof(β.val), conj(dot(C.dval, Cval))) + else + ntuple(i -> _project(typeof(β.val), conj(dot(C.dval[i], Cval))), Val(N)) + end + else + nothing + end + + αc = conj(α.val) + βc = conj(β.val) + ntuple(Val(N)) do i + Base.@_inline_meta + dC = _bget(C.dval, Val(N), i) + if !(A isa Const) + _bget(A.dval, Val(N), i) .+= αc .* (dC * adjoint(Bval)) + end + if !(B isa Const) + _bget(B.dval, Val(N), i) .+= αc .* (adjoint(Aval) * dC) + end + dC .*= βc + nothing + end + else + dα = !(α isa Const) ? (N == 1 ? zero(α.val) : ntuple(Returns(zero(α.val)), Val(N))) : nothing + dβ = !(β isa Const) ? (N == 1 ? zero(β.val) : ntuple(Returns(zero(β.val)), Val(N))) : nothing + end + + return (nothing, nothing, nothing, dα, dβ) +end + +#= +A * B: allocating matmul (matrix or matrix-vector) + +Unlike the in-place `mul!` rule above, `*` allocates its result, so the rule +owns the output shadow: `augmented_primal` returns a zeroed shadow and +`reverse` zeroes it again after reading. Without this, the freshly-allocated +result's shadow is uninitialized and downstream `+=` accumulates onto garbage. + +JVP: dY = dA·B + A·dB +Pullback: dA += dY·B' + dB += A'·dY +=# + +# dY_i = dA_i·B + A·dB_i (factored out so inference sees a concrete array type) +@inline function _matmul_jvp(A::Annotation, B::Annotation, ::Val{N}, i::Int) where {N} + if A isa Const + return A.val * _bget(B.dval, Val(N), i) + elseif B isa Const + return _bget(A.dval, Val(N), i) * B.val + else + return _bget(A.dval, Val(N), i) * B.val .+ A.val * _bget(B.dval, Val(N), i) + end +end + +function EnzymeRules.forward( + config::FwdConfig, + func::Const{typeof(*)}, + RT::Type{<:Annotation}, + A::Annotation{<:MaybeWrappedGPU}, + B::Annotation{<:MaybeWrappedGPU}, + ) + if RT <: Const + return needs_primal(config) ? A.val * B.val : nothing + end + + N = width(config) + if N == 1 + dY = _matmul_jvp(A, B, Val(1), 1) + return RT <: DuplicatedNoNeed ? dY : Duplicated(A.val * B.val, dY) + else + dY = ntuple(i -> _matmul_jvp(A, B, Val(N), i), Val(N)) + return RT <: BatchDuplicatedNoNeed ? dY : BatchDuplicated(A.val * B.val, dY) + end +end + +function EnzymeRules.augmented_primal( + config::RevConfig, + func::Const{typeof(*)}, + ::Type{RT}, + A::Annotation{<:MaybeWrappedGPU}, + B::Annotation{<:MaybeWrappedGPU}, + ) where {RT} + Y = A.val * B.val + primal = needs_primal(config) ? Y : nothing + + N = width(config) + dY = if RT <: Duplicated || RT <: DuplicatedNoNeed + zero(Y) + elseif RT <: BatchDuplicated || RT <: BatchDuplicatedNoNeed + ntuple(_ -> zero(Y), Val(N)) + else + nothing + end + + cache_A = (overwritten(config)[2] && !(B isa Const)) ? copy(A.val) : A.val + cache_B = (overwritten(config)[3] && !(A isa Const)) ? copy(B.val) : B.val + + return AugmentedReturn(primal, dY, (dY, cache_A, cache_B)) +end + +function EnzymeRules.reverse( + config::RevConfig, + func::Const{typeof(*)}, + ::Type{RT}, + tape, + A::Annotation{<:MaybeWrappedGPU}, + B::Annotation{<:MaybeWrappedGPU}, + ) where {RT} + dY, cache_A, cache_B = tape + if dY !== nothing + N = width(config) + ntuple(Val(N)) do i + Base.@_inline_meta + dYi = _bget(dY, Val(N), i) + if !(A isa Const) + _bget(A.dval, Val(N), i) .+= dYi * adjoint(cache_B) + end + if !(B isa Const) + _bget(B.dval, Val(N), i) .+= adjoint(cache_A) * dYi + end + fill!(dYi, zero(eltype(dYi))) + nothing + end + end + return (nothing, nothing) +end + +#= +dot(a, b): scalar inner product + +JVP: dr = ⟨da, b⟩ + ⟨a, db⟩ +Pullback: da += dr̄·b + db += dr̄·a (real convention; targets real GPU workloads) +=# + +function EnzymeRules.forward( + config::FwdConfig, + func::Const{typeof(dot)}, + RT::Type{<:Annotation}, + a::Annotation{<:AbstractGPUArray}, + b::Annotation{<:AbstractGPUArray}, + ) + if needs_shadow(config) + N = width(config) + z = zero(promote_type(eltype(a.val), eltype(b.val))) + if N == 1 + dr = (a isa Const ? z : dot(a.dval, b.val)) + (b isa Const ? z : dot(a.val, b.dval)) + return needs_primal(config) ? Duplicated(dot(a.val, b.val), dr) : dr + else + dr = ntuple( + i -> (a isa Const ? z : dot(_bget(a.dval, Val(N), i), b.val)) + + (b isa Const ? z : dot(a.val, _bget(b.dval, Val(N), i))), + Val(N), + ) + return needs_primal(config) ? BatchDuplicated(dot(a.val, b.val), dr) : dr + end + elseif needs_primal(config) + return dot(a.val, b.val) + else + return nothing + end +end + +function EnzymeRules.augmented_primal( + config::RevConfig, + func::Const{typeof(dot)}, + ::Type, + a::Annotation{<:AbstractGPUArray}, + b::Annotation{<:AbstractGPUArray}, + ) + primal = needs_primal(config) ? dot(a.val, b.val) : nothing + cache_a = (overwritten(config)[2] && !(b isa Const)) ? copy(a.val) : nothing + cache_b = (overwritten(config)[3] && !(a isa Const)) ? copy(b.val) : nothing + return AugmentedReturn(primal, nothing, (cache_a, cache_b)) +end + +function EnzymeRules.reverse( + config::RevConfig, + func::Const{typeof(dot)}, + dret, + tape, + a::Annotation{<:AbstractGPUArray}, + b::Annotation{<:AbstractGPUArray}, + ) + if !(dret isa Const) + cache_a, cache_b = tape + av = cache_a !== nothing ? cache_a : a.val + bv = cache_b !== nothing ? cache_b : b.val + N = width(config) + ntuple(Val(N)) do i + Base.@_inline_meta + dr = _bget(dret.val, Val(N), i) + if !(a isa Const) + _bget(a.dval, Val(N), i) .+= dr .* bv + end + if !(b isa Const) + _bget(b.dval, Val(N), i) .+= dr .* av + end + nothing + end + end + return (nothing, nothing) +end + +#= +sum(x): scalar reduction + +JVP: dr = sum(dx) +Pullback: dx += dr̄ (scalar broadcast over every element) +=# + +function EnzymeRules.forward( + config::FwdConfig, + func::Const{typeof(sum)}, + RT::Type{<:Annotation}, + x::Annotation{<:AbstractGPUArray}, + ) + if needs_shadow(config) + N = width(config) + if N == 1 + dr = x isa Const ? zero(eltype(x.val)) : sum(x.dval) + return needs_primal(config) ? Duplicated(sum(x.val), dr) : dr + else + dr = ntuple(i -> x isa Const ? zero(eltype(x.val)) : sum(_bget(x.dval, Val(N), i)), Val(N)) + return needs_primal(config) ? BatchDuplicated(sum(x.val), dr) : dr + end + elseif needs_primal(config) + return sum(x.val) + else + return nothing + end +end + +function EnzymeRules.augmented_primal( + config::RevConfig, + func::Const{typeof(sum)}, + ::Type, + x::Annotation{<:AbstractGPUArray}, + ) + primal = needs_primal(config) ? sum(x.val) : nothing + return AugmentedReturn(primal, nothing, nothing) +end + +function EnzymeRules.reverse( + config::RevConfig, + func::Const{typeof(sum)}, + dret, + tape, + x::Annotation{<:AbstractGPUArray}, + ) + if !(dret isa Const) && !(x isa Const) + N = width(config) + ntuple(Val(N)) do i + Base.@_inline_meta + _bget(x.dval, Val(N), i) .+= _bget(dret.val, Val(N), i) + nothing + end + end + return (nothing,) +end + end # module diff --git a/test/cuda.jl b/test/cuda.jl index d3bfb8e2b0..f28c75867e 100644 --- a/test/cuda.jl +++ b/test/cuda.jl @@ -1,6 +1,7 @@ using CUDA using Enzyme using Test +using LinearAlgebra: mul!, dot function mul_kernel(A) i = threadIdx().x @@ -191,5 +192,87 @@ end dA2 .= 3 Enzyme.autodiff(Reverse, square!, BatchDuplicated(A, (dA, dA2))) @test all(dA .≈ (2:2:64)) - @test all(dA2 .≈ 3*(2:2:64)) -end \ No newline at end of file + @test all(dA2 .≈ 3 * (2:2:64)) +end + +# Rules from EnzymeGPUArraysCoreExt for dense matmul / dot / sum on GPU arrays. +@testset "GPUArrays linalg rules" begin + cu(x) = CuArray(x) + + @testset "matmul reverse" begin + A0 = randn(Float32, 3, 4) + B0 = randn(Float32, 4, 2) + dA = cu(zero(A0)) + dB = cu(zero(B0)) + Enzyme.autodiff(Reverse, (A, B) -> sum(A * B), Active, Duplicated(cu(A0), dA), Duplicated(cu(B0), dB)) + @test Array(dA) ≈ ones(Float32, 3, 2) * B0' + @test Array(dB) ≈ A0' * ones(Float32, 3, 2) + end + + @testset "composed transpose matmul" begin + X0 = randn(Float32, 6, 3) + β0 = randn(Float32, 3) + dX = cu(zero(X0)) + dβ = cu(zero(β0)) + Enzyme.autodiff(Reverse, (X, β) -> sum(transpose(X) * (X * β)), Active, Duplicated(cu(X0), dX), Duplicated(cu(β0), dβ)) + ϵ = 1.0f-3 + fdβ = map(eachindex(β0)) do i + βp = copy(β0); βp[i] += ϵ + βm = copy(β0); βm[i] -= ϵ + (sum(transpose(X0) * (X0 * βp)) - sum(transpose(X0) * (X0 * βm))) / (2ϵ) + end + @test Array(dβ) ≈ fdβ rtol = 1.0f-2 + end + + @testset "mul! reverse (preallocated)" begin + A0 = randn(Float32, 3, 4) + B0 = randn(Float32, 4, 2) + dA = cu(zero(A0)) + dB = cu(zero(B0)) + function loss(C, A, B) + mul!(C, A, B) + return sum(C) + end + C = cu(zeros(Float32, 3, 2)) + dC = cu(zeros(Float32, 3, 2)) + Enzyme.autodiff(Reverse, loss, Active, Duplicated(C, dC), Duplicated(cu(A0), dA), Duplicated(cu(B0), dB)) + @test Array(dA) ≈ ones(Float32, 3, 2) * B0' + @test Array(dB) ≈ A0' * ones(Float32, 3, 2) + end + + @testset "dot reverse" begin + a0 = randn(Float32, 8) + b0 = randn(Float32, 8) + da = cu(zero(a0)) + db = cu(zero(b0)) + Enzyme.autodiff(Reverse, (a, b) -> dot(a, b), Active, Duplicated(cu(a0), da), Duplicated(cu(b0), db)) + @test Array(da) ≈ b0 + @test Array(db) ≈ a0 + end + + @testset "sum reverse" begin + x0 = randn(Float32, 10) + dx = cu(zero(x0)) + Enzyme.autodiff(Reverse, sum, Active, Duplicated(cu(x0), dx)) + @test all(Array(dx) .≈ 1) + end + + @testset "sum(A .* B) broadcast reduction" begin + A0 = randn(Float32, 4, 5) + B0 = randn(Float32, 4, 5) + dA = cu(zero(A0)) + dB = cu(zero(B0)) + Enzyme.autodiff(Reverse, (A, B) -> sum(A .* B), Active, Duplicated(cu(A0), dA), Duplicated(cu(B0), dB)) + @test Array(dA) ≈ B0 + @test Array(dB) ≈ A0 + end + + @testset "matmul forward" begin + A0 = randn(Float32, 3, 4) + B0 = randn(Float32, 4, 2) + dA = randn(Float32, 3, 4) + dB = randn(Float32, 4, 2) + out = Enzyme.autodiff(Forward, (A, B) -> A * B, Duplicated, Duplicated(cu(A0), cu(dA)), Duplicated(cu(B0), cu(dB))) + @test Array(out[1]) ≈ dA * B0 + A0 * dB + end +end diff --git a/test/ext/jlarrays.jl b/test/ext/jlarrays.jl index 760000ba75..e5ac218413 100644 --- a/test/ext/jlarrays.jl +++ b/test/ext/jlarrays.jl @@ -1,4 +1,5 @@ using Enzyme, Test, JLArrays +using LinearAlgebra: mul!, dot function jlres(x) 2 * collect(x) @@ -9,3 +10,79 @@ end # Enzyme.jacobian(Forward, jlres, JLArray([3.0, 5.0])) # Enzyme.jacobian(Reverse, jlres, JLArray([3.0, 5.0])) end + +# AbstractGPUArray linear-algebra rules (matmul / dot / sum). JLArray is a +# CPU-backed `AbstractGPUArray`, so it exercises exactly the dispatch the +# CUDA/AMDGPU/Metal backends hit, without needing a device. +@testset "GPUArrays linalg rules" begin + jl(x) = JLArray(x) + + @testset "matmul reverse ($m×$k × $k×$n)" for (m, k, n) in ((3, 4, 2), (5, 5, 1)) + A0 = randn(m, k) + B0 = randn(k, n) + + # loss = sum(A*B); analytic grads dA = ones*B', dB = A'*ones + f(A, B) = sum(A * B) + dA = jl(zero(A0)) + dB = jl(zero(B0)) + Enzyme.autodiff(Reverse, f, Active, Duplicated(jl(A0), dA), Duplicated(jl(B0), dB)) + + ones_mn = ones(m, n) + @test collect(dA) ≈ ones_mn * B0' + @test collect(dB) ≈ A0' * ones_mn + end + + @testset "matmul reverse with transpose" begin + # mirrors pmcmc_matmul(transpose(X), X*β): a composed/wrapped matmul + X0 = randn(6, 3) + β0 = randn(3) + g(X, β) = sum(transpose(X) * (X * β)) + dX = jl(zero(X0)) + dβ = jl(zero(β0)) + Enzyme.autodiff(Reverse, g, Active, Duplicated(jl(X0), dX), Duplicated(jl(β0), dβ)) + + # finite-difference check against the CPU primal + gcpu(X, β) = sum(transpose(X) * (X * β)) + ϵ = 1.0e-6 + fdβ = map(eachindex(β0)) do i + βp = copy(β0); βp[i] += ϵ + βm = copy(β0); βm[i] -= ϵ + (gcpu(X0, βp) - gcpu(X0, βm)) / (2ϵ) + end + @test collect(dβ) ≈ fdβ rtol = 1.0e-4 + end + + @testset "dot reverse" begin + a0 = randn(8) + b0 = randn(8) + h(a, b) = dot(a, b) + da = jl(zero(a0)) + db = jl(zero(b0)) + Enzyme.autodiff(Reverse, h, Active, Duplicated(jl(a0), da), Duplicated(jl(b0), db)) + @test collect(da) ≈ b0 + @test collect(db) ≈ a0 + end + + @testset "sum reverse" begin + x0 = randn(10) + da = jl(zero(x0)) + Enzyme.autodiff(Reverse, sum, Active, Duplicated(jl(x0), da)) + @test all(collect(da) .≈ 1) + # Note: `sum(A .* B)` (reduction over a broadcast) also relies on this + # rule, but JLArrays cannot differentiate the broadcast kernel itself, + # so that pattern is covered in the CUDA test instead. + end + + @testset "matmul forward" begin + A0 = randn(3, 4) + B0 = randn(4, 2) + dA = randn(3, 4) + dB = randn(4, 2) + # d(A*B) = dA*B + A*dB + out = Enzyme.autodiff( + Forward, (A, B) -> A * B, Duplicated, + Duplicated(jl(A0), jl(dA)), Duplicated(jl(B0), jl(dB)), + ) + @test collect(out[1]) ≈ dA * B0 + A0 * dB + end +end From 48411be65bd5a992b7e5e69e0f9794a7e130cba4 Mon Sep 17 00:00:00 2001 From: Ryan Senne Date: Sun, 7 Jun 2026 19:27:16 -0400 Subject: [PATCH 2/3] Constrain Operands to GPU array types --- ext/EnzymeGPUArraysCoreExt.jl | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/ext/EnzymeGPUArraysCoreExt.jl b/ext/EnzymeGPUArraysCoreExt.jl index b9548eb249..486c148910 100644 --- a/ext/EnzymeGPUArraysCoreExt.jl +++ b/ext/EnzymeGPUArraysCoreExt.jl @@ -78,8 +78,8 @@ function EnzymeRules.forward( func::Const{typeof(mul!)}, RT::Type{<:Annotation}, C::Annotation{<:AbstractGPUArray}, - A::Annotation, - B::Annotation, + A::Annotation{<:MaybeWrappedGPU}, + B::Annotation{<:MaybeWrappedGPU}, α::Annotation{<:Number}, β::Annotation{<:Number}, ) @@ -126,8 +126,8 @@ function EnzymeRules.augmented_primal( func::Const{typeof(mul!)}, ::Type{RT}, C::Annotation{<:AbstractGPUArray}, - A::Annotation, - B::Annotation, + A::Annotation{<:MaybeWrappedGPU}, + B::Annotation{<:MaybeWrappedGPU}, α::Annotation{<:Number}, β::Annotation{<:Number}, ) where {RT} @@ -153,8 +153,8 @@ function EnzymeRules.reverse( ::Type{RT}, tape, C::Annotation{<:AbstractGPUArray}, - A::Annotation, - B::Annotation, + A::Annotation{<:MaybeWrappedGPU}, + B::Annotation{<:MaybeWrappedGPU}, α::Annotation{<:Number}, β::Annotation{<:Number}, ) where {RT} From a1733445be5e31b09c070edac17aeb590cd9cb06 Mon Sep 17 00:00:00 2001 From: Ryan Senne Date: Sun, 7 Jun 2026 20:18:30 -0400 Subject: [PATCH 3/3] Support structured matmul operands (triangular/symmetric/hermitian) Extend the GPUArrays matmul rules to cover UpperTriangular, LowerTriangular, Symmetric, and Hermitian operands. The primal already dispatches to the specialized BLAS kernel (trmm/symm/hemm); the reverse pass now projects the dense cotangent onto each wrapper's stored entries via _accumulate_operand!. UnitTriangular is excluded since its fixed diagonal cannot carry a tangent. Verified against finite differences on JLArrays and CUDA. --- ext/EnzymeGPUArraysCoreExt.jl | 67 +++++++++++++++++++++++++++++++---- test/cuda.jl | 30 +++++++++++++++- test/ext/jlarrays.jl | 52 ++++++++++++++++++++++++++- 3 files changed, 140 insertions(+), 9 deletions(-) diff --git a/ext/EnzymeGPUArraysCoreExt.jl b/ext/EnzymeGPUArraysCoreExt.jl index 486c148910..890eacbebd 100644 --- a/ext/EnzymeGPUArraysCoreExt.jl +++ b/ext/EnzymeGPUArraysCoreExt.jl @@ -2,7 +2,8 @@ module EnzymeGPUArraysCoreExt using GPUArraysCore using Enzyme -using LinearAlgebra: LinearAlgebra, mul!, dot, Transpose, Adjoint, adjoint +using LinearAlgebra: LinearAlgebra, mul!, dot, Transpose, Adjoint, adjoint, + UpperTriangular, LowerTriangular, Symmetric, Hermitian, triu, tril, diag, diagind using Enzyme.EnzymeCore: EnzymeCore using Enzyme.EnzymeCore.EnzymeRules: EnzymeRules, @@ -54,14 +55,66 @@ end _project(::Type{<:Real}, x) = real(x) _project(::Type, x) = x -# A GPU array or a (lazy) transpose/adjoint of one; the operand types that show -# up in `A * B` matmuls, including `transpose(X) * y`. +# A GPU array, or a structured/lazy wrapper of one; the operand types that show +# up in `A * B` matmuls (`transpose(X) * y`, `Symmetric(A) * x`, ...). CUBLAS / +# rocBLAS dispatch these to specialized kernels (trmm/symm/hemm), and the reverse +# pass projects the cotangent back onto each wrapper's stored entries (see +# `_accumulate_operand!`). UnitTriangular is intentionally excluded: its diagonal +# is structurally fixed, so it cannot represent a tangent/cotangent. const MaybeWrappedGPU = Union{ AbstractGPUArray, Transpose{<:Any, <:AbstractGPUArray}, Adjoint{<:Any, <:AbstractGPUArray}, + UpperTriangular{<:Any, <:AbstractGPUArray}, + LowerTriangular{<:Any, <:AbstractGPUArray}, + Symmetric{<:Any, <:AbstractGPUArray}, + Hermitian{<:Any, <:AbstractGPUArray}, } +#= +Accumulate `factor .* G` into the shadow `s` of a matmul operand, where `G` is +the dense cotangent (`dY·B'` or `A'·dY`). For dense / transpose / adjoint shadows +the cotangent is added directly. For structured shadows only the stored entries +are free parameters, so `G` is projected onto that structure: + + UpperTriangular : keep `triu(G)` + LowerTriangular : keep `tril(G)` + Symmetric(uplo) : off-diagonal (i,j) collects `G[i,j] + G[j,i]`, diagonal `G[i,i]` + Hermitian(uplo) : same with conjugation; the diagonal is real +=# +@inline function _accumulate_operand!(s, G, factor) + s .+= factor .* G + return nothing +end + +function _accumulate_operand!(s::UpperTriangular, G, factor) + parent(s) .+= factor .* triu(G) + return nothing +end + +function _accumulate_operand!(s::LowerTriangular, G, factor) + parent(s) .+= factor .* tril(G) + return nothing +end + +function _accumulate_operand!(s::Symmetric, G, factor) + p = parent(s) + H = G .+ transpose(G) + p .+= factor .* (s.uplo == 'U' ? triu(H, 1) : tril(H, -1)) + dg = view(p, diagind(p)) + dg .+= factor .* diag(G) + return nothing +end + +function _accumulate_operand!(s::Hermitian, G, factor) + p = parent(s) + H = G .+ adjoint(G) + p .+= factor .* (s.uplo == 'U' ? triu(H, 1) : tril(H, -1)) + dg = view(p, diagind(p)) + dg .+= factor .* real.(diag(G)) + return nothing +end + #= mul!(C, A, B, α, β): C = α·A·B + β·C₀ @@ -190,10 +243,10 @@ function EnzymeRules.reverse( Base.@_inline_meta dC = _bget(C.dval, Val(N), i) if !(A isa Const) - _bget(A.dval, Val(N), i) .+= αc .* (dC * adjoint(Bval)) + _accumulate_operand!(_bget(A.dval, Val(N), i), dC * adjoint(Bval), αc) end if !(B isa Const) - _bget(B.dval, Val(N), i) .+= αc .* (adjoint(Aval) * dC) + _accumulate_operand!(_bget(B.dval, Val(N), i), adjoint(Aval) * dC, αc) end dC .*= βc nothing @@ -291,10 +344,10 @@ function EnzymeRules.reverse( Base.@_inline_meta dYi = _bget(dY, Val(N), i) if !(A isa Const) - _bget(A.dval, Val(N), i) .+= dYi * adjoint(cache_B) + _accumulate_operand!(_bget(A.dval, Val(N), i), dYi * adjoint(cache_B), true) end if !(B isa Const) - _bget(B.dval, Val(N), i) .+= adjoint(cache_A) * dYi + _accumulate_operand!(_bget(B.dval, Val(N), i), adjoint(cache_A) * dYi, true) end fill!(dYi, zero(eltype(dYi))) nothing diff --git a/test/cuda.jl b/test/cuda.jl index f28c75867e..22baca81ad 100644 --- a/test/cuda.jl +++ b/test/cuda.jl @@ -1,7 +1,7 @@ using CUDA using Enzyme using Test -using LinearAlgebra: mul!, dot +using LinearAlgebra: mul!, dot, UpperTriangular, LowerTriangular, Symmetric, Hermitian function mul_kernel(A) i = threadIdx().x @@ -275,4 +275,32 @@ end out = Enzyme.autodiff(Forward, (A, B) -> A * B, Duplicated, Duplicated(cu(A0), cu(dA)), Duplicated(cu(B0), cu(dB))) @test Array(out[1]) ≈ dA * B0 + A0 * dB end + + # Structured operands: primal uses the specialized BLAS kernel; reverse + # projects the cotangent onto the wrapper's stored entries. Ground truth is + # central finite differences over the underlying data. + @testset "structured operand: $name" for (name, wrap) in ( + ("UpperTriangular", UpperTriangular), + ("LowerTriangular", LowerTriangular), + ("Symmetric(:U)", X -> Symmetric(X, :U)), + ("Symmetric(:L)", X -> Symmetric(X, :L)), + ("Hermitian(:U)", X -> Hermitian(X, :U)), + ) + n = 4 + X0 = randn(Float32, n, n) + B0 = randn(Float32, n, 3) + dX = cu(zero(X0)) + Enzyme.autodiff( + Reverse, (A, B) -> sum(A * B), Active, + Duplicated(wrap(cu(X0)), wrap(dX)), Const(cu(B0)), + ) + ϵ = 1.0f-3 + fd = zero(X0) + for idx in eachindex(X0) + Xp = copy(X0); Xp[idx] += ϵ + Xm = copy(X0); Xm[idx] -= ϵ + fd[idx] = (sum(wrap(Xp) * B0) - sum(wrap(Xm) * B0)) / (2ϵ) + end + @test Array(dX) ≈ fd rtol = 1.0f-2 + end end diff --git a/test/ext/jlarrays.jl b/test/ext/jlarrays.jl index e5ac218413..555ebbbf76 100644 --- a/test/ext/jlarrays.jl +++ b/test/ext/jlarrays.jl @@ -1,5 +1,5 @@ using Enzyme, Test, JLArrays -using LinearAlgebra: mul!, dot +using LinearAlgebra: mul!, dot, UpperTriangular, LowerTriangular, Symmetric, Hermitian function jlres(x) 2 * collect(x) @@ -85,4 +85,54 @@ end ) @test collect(out[1]) ≈ dA * B0 + A0 * dB end + + # Structured operands (triangular / symmetric / hermitian). The primal goes + # through the specialized BLAS kernel; the reverse must project the cotangent + # onto the wrapper's stored entries. Ground truth is central finite + # differences over the underlying data (non-stored entries must stay 0). + @testset "structured operand: $name" for (name, wrap) in ( + ("UpperTriangular", UpperTriangular), + ("LowerTriangular", LowerTriangular), + ("Symmetric(:U)", X -> Symmetric(X, :U)), + ("Symmetric(:L)", X -> Symmetric(X, :L)), + ("Hermitian(:U)", X -> Hermitian(X, :U)), + ) + n = 4 + X0 = randn(n, n) + B0 = randn(n, 3) + + dX = jl(zero(X0)) + Enzyme.autodiff( + Reverse, (A, B) -> sum(A * B), Active, + Duplicated(wrap(jl(X0)), wrap(dX)), Const(jl(B0)), + ) + + ϵ = 1.0e-6 + fd = zero(X0) + for idx in eachindex(X0) + Xp = copy(X0); Xp[idx] += ϵ + Xm = copy(X0); Xm[idx] -= ϵ + fd[idx] = (sum(wrap(Xp) * B0) - sum(wrap(Xm) * B0)) / (2ϵ) + end + @test collect(dX) ≈ fd rtol = 1.0e-5 + end + + @testset "structured operand on the right (B = UpperTriangular)" begin + m, n = 3, 4 + A0 = randn(m, n) + X0 = randn(n, n) + dX = jl(zero(X0)) + Enzyme.autodiff( + Reverse, (A, B) -> sum(A * B), Active, + Const(jl(A0)), Duplicated(UpperTriangular(jl(X0)), UpperTriangular(dX)), + ) + ϵ = 1.0e-6 + fd = zero(X0) + for idx in eachindex(X0) + Xp = copy(X0); Xp[idx] += ϵ + Xm = copy(X0); Xm[idx] -= ϵ + fd[idx] = (sum(A0 * UpperTriangular(Xp)) - sum(A0 * UpperTriangular(Xm))) / (2ϵ) + end + @test collect(dX) ≈ fd rtol = 1.0e-5 + end end