Add rules for BLAS.dot, BLAS.dotc, and BLAS.dotu

src/Enzyme.jl

@@ -45,6 +45,7 @@ include("gradientutils.jl")
 include("utils.jl")
 include("compiler.jl")
 include("internal_rules.jl")
+include("rules/LinearAlgebra/blas.jl")
 
 import .Compiler: CompilationException
 

src/rules/LinearAlgebra/blas.jl

@@ -0,0 +1,149 @@
+module BLASRules
+
+using ..Enzyme
+using LinearAlgebra.BLAS
+
+const ConstOrDuplicated{T} = Union{Const{T},Duplicated{T}}
+const ConstOrBatchDuplicated{T} = Union{ConstOrDuplicated{T},BatchDuplicated{T}}
+
+_safe_similar(x::AbstractArray, n::Integer) = similar(x, n)
+_safe_similar(x::Ptr, n::Integer) = Array{eltype(x)}(undef, n)
+
+function _strided_tape(n::Integer, x::Union{AbstractArray,Ptr}, incx::Integer)
+    xtape = _safe_similar(x, n)
+    BLAS.blascopy!(n, x, incx, xtape, 1)
+    increment = 1
+    return xtape, increment
+end
+
+function _strided_range(n, x, incx)
+    r = range(1; step=abs(incx), length=n)
+    incx < 0 && return reverse(r)
+    return r
+end
+
+function _strided_view(n::Integer, x::AbstractArray, incx::Integer)
+    ind = _strided_range(n, x, incx)
+    return view(x, ind)
+end
+function _strided_view(n::Integer, x::Ptr, incx::Integer)
+    ind = _strided_range(n, x, incx)
+    dim = abs(last(ind) - first(ind)) + 1
+    y = Base.unsafe_wrap(Array, x, dim)
+    return view(y, ind)
+end
+
+function _maybe_primal_shadow(config, func, args)
+    needs_primal = EnzymeRules.needs_primal(config)
+    needs_shadow = EnzymeRules.needs_shadow(config)
+    if needs_primal || needs_shadow
+        r = func(args...)
+    else
+        r = nothing
+    end
+    primal = needs_primal ? r : nothing
+    shadow = needs_shadow ? zero(r) : nothing
+    return primal, shadow
+end
+
+_map_tuple(f, xs::Tuple...) = map(f, xs...)
+_map_tuple(f, xs...) = f(xs...)
+
+# axpy!(a, conj.(x), y)
+function _aconjxpy!(n, a, x, incx, y, incy)
+    xview = _strided_view(n, x, incx)
+    yview = _strided_view(n, y, incy)
+    yview .+= a .* conj.(xview)
+    return y
+end
+
+for (fname, Ttype) in ((:dot, :BlasReal), (:dotu, :BlasComplex), (:dotc, :BlasComplex))
+    @eval begin
+        function EnzymeRules.forward(
+            func::Const{typeof(BLAS.$fname)},
+            RT::Type{
+                <:Union{
+                    Const,DuplicatedNoNeed,Duplicated,BatchDuplicatedNoNeed,BatchDuplicated
+                },
+            },
+            n::Const{<:Integer},
+            X::ConstOrBatchDuplicated{<:Union{Ptr{T},AbstractArray{T}}},
+            incx::Const{<:Integer},
+            Y::ConstOrBatchDuplicated{<:Union{Ptr{T},AbstractArray{T}}},
+            incy::Const{<:Integer},
+        ) where {T<:BLAS.$Ttype}
+            RT <: Const && return func.val(n.val, X.val, incx.val, Y.val, incy.val)
+
+            dval = if !(X isa Const) && !(Y isa Const)
+                _map_tuple(X.dval, Y.dval) do dX, dY
+                    func.val(n.val, dX, incx.val, Y.val, incy.val) +
+                    func.val(n.val, X.val, incx.val, dY, incy.val)
+                end
+            elseif !(X isa Const)
+                _map_tuple(dX -> func.val(n.val, dX, incx.val, Y.val, incy.val), X.dval)
+            elseif !(Y isa Const)
+                _map_tuple(dY -> func.val(n.val, X.val, incx.val, dY, incy.val), Y.dval)
+            else
+                zero(T)
+            end
+
+            if RT <: Union{DuplicatedNoNeed,BatchDuplicatedNoNeed}
+                return dval
+            else
+                val = func.val(n.val, X.val, incx.val, Y.val, incy.val)
+                return RT(val, dval)
+            end
+        end
+
+        function EnzymeRules.augmented_primal(
+            config::EnzymeRules.ConfigWidth{1},
+            func::Const{typeof(BLAS.$fname)},
+            RT::Type{<:Union{Const,Active}},
+            n::Const{<:Integer},
+            X::ConstOrDuplicated{<:Union{Ptr{T},AbstractArray{T}}},
+            incx::Const{<:Integer},
+            Y::ConstOrDuplicated{<:Union{Ptr{T},AbstractArray{T}}},
+            incy::Const{<:Integer},
+        ) where {T<:BLAS.$Ttype}
+            primal, shadow = _maybe_primal_shadow(
+                config, func.val, (n.val, X.val, incx.val, Y.val, incy.val)
+            )
+
+            # build tape
+            _, _, Xow, _, Yow = EnzymeRules.overwritten(config)
+            tape_X = !(RT <: Const) && !(Y isa Const) && Xow
+            tape_Y = !(RT <: Const) && !(X isa Const) && Yow
+            Xtape = tape_X ? _strided_tape(n.val, X.val, incx.val) : (X.val, incx.val)
+            Ytape = tape_Y ? _strided_tape(n.val, Y.val, incy.val) : (Y.val, incy.val)
+            tape = (Xtape, Ytape)
+
+            return EnzymeRules.AugmentedReturn(primal, shadow, tape)
+        end
+
+        function EnzymeRules.reverse(
+            config::EnzymeRules.ConfigWidth{1},
+            fun::Const{typeof(BLAS.$fname)},
+            dret::Union{Active,Type{<:Const}},
+            tape,
+            n::Const{<:Integer},
+            X::ConstOrDuplicated{<:Union{Ptr{T},AbstractArray{T}}},
+            incx::Const{<:Integer},
+            Y::ConstOrDuplicated{<:Union{Ptr{T},AbstractArray{T}}},
+            incy::Const{<:Integer},
+        ) where {T<:BLAS.$Ttype}
+            ret = (nothing, nothing, nothing, nothing, nothing)
+            dret isa Type{<:Const} && return ret
+
+            (Xval, incxval), (Yval, incyval) = tape
+
+            atransxpy! = fun.val === BLAS.dotu ? _aconjxpy! : BLAS.axpy!
+            dval_X = fun.val === BLAS.dotu ? dret.val : conj(dret.val)
+            X isa Const || atransxpy!(n.val, dval_X, Yval, incyval, X.dval, incx.val)
+            Y isa Const || atransxpy!(n.val, dret.val, Xval, incxval, Y.dval, incy.val)
+
+            return ret
+        end
+    end
+end
+
+end  # module

test/rules/LinearAlgebra/blas.jl

@@ -0,0 +1,190 @@
+using Enzyme
+using FiniteDifferences
+using LinearAlgebra
+using Test
+
+@testset "BLAS rules" begin
+    fdm = central_fdm(5, 1)
+    RTs = (Float32, Float64)
+    RCs = (ComplexF32, ComplexF64)
+    n = 10
+
+    @testset for fun in (BLAS.dot, BLAS.dotu, BLAS.dotc)
+        @testset "forward" begin
+            @testset for Tret in (Const, Duplicated, DuplicatedNoNeed),
+                Tx in (Const, Duplicated),
+                Ty in (Const, Duplicated),
+                pfun in (identity, pointer),
+                T in (fun == BLAS.dot ? RTs : RCs),
+                (sz, inc) in ((10, 1), ((2, 20), -2))
+
+                Tx <: Const && Ty <: Const && !(Tret <: Const) && continue
+
+                x, ∂x = ntuple(_ -> randn(T, sz), 2)
+                y, ∂y = ntuple(_ -> randn(T, sz), 2)
+
+                x_annot = Tx <: Const ? Const(pfun(x)) : Duplicated(pfun(x), pfun(∂x))
+                y_annot = Ty <: Const ? Const(pfun(y)) : Duplicated(pfun(y), pfun(∂y))
+
+                vexp = fun(n, x, inc, y, inc)
+                dexp = FiniteDifferences.jvp(
+                    fdm,
+                    (x, y) -> fun(n, x, inc, y, inc),
+                    Tx <: Const ? (x, zero(x)) : (x, ∂x),
+                    Ty <: Const ? (y, zero(y)) : (y, ∂y),
+                )[1]
+                ret = autodiff(
+                    Forward, fun, Tret, Const(n), x_annot, Const(inc), y_annot, Const(inc)
+                )
+
+                Tret <: Const && @test ret === ()
+                if Tret <: Duplicated
+                    v, d = ret
+                    @test v ≈ vexp
+                    @test d ≈ dexp
+                elseif Tret <: DuplicatedNoNeed
+                    @test only(ret) ≈ dexp
+                end
+
+                if pfun === identity && sz == n && inc == 1
+                    @testset "consistency of 2-arg version" begin
+                        ret2 = autodiff(Forward, fun, Tret, x_annot, y_annot)
+                        @test ret2 == ret
+                    end
+                end
+            end
+
+            @testset for Tret in (BatchDuplicated, BatchDuplicatedNoNeed),
+                T in (fun == BLAS.dot ? RTs : RCs)
+
+                batch_size = 3
+                inc = 1
+                x = randn(T, n)
+                y = randn(T, n)
+                ∂xs = ntuple(_ -> randn(T, n), batch_size)
+                ∂ys = ntuple(_ -> randn(T, n), batch_size)
+                vexp = fun(n, x, inc, y, inc)
+                dexp = map(∂xs, ∂ys) do ∂x, ∂y
+                    FiniteDifferences.jvp(
+                        fdm, (x, y) -> fun(n, x, inc, y, inc), (x, ∂x), (y, ∂y)
+                    )[1]
+                end
+                ret = autodiff(
+                    Forward,
+                    fun,
+                    Tret,
+                    Const(n),
+                    BatchDuplicated(x, ∂xs),
+                    Const(inc),
+                    BatchDuplicated(y, ∂ys),
+                    Const(inc),
+                )
+                if Tret <: BatchDuplicated
+                    v, ds = ret
+                    @test v ≈ vexp
+                else
+                    ds = only(ret)
+                end
+                @test all(map(≈, values(ds), dexp))
+            end
+        end
+
+        @testset "reverse" begin
+            function fun_overwrite!(n, x, incx, y, incy)
+                d = fun(n, x, incx, y, incy)
+                x[1] = 0
+                y[1] = 0
+                return d
+            end
+            function fun_overwrite!(x, y)
+                d = fun(x, y)
+                x[1] = 0
+                y[1] = 0
+                return d
+            end
+
+            @testset for Tret in (Const, Active),
+                Tx in (Const, Duplicated),
+                Ty in (Const, Duplicated),
+                pfun in (identity, pointer),
+                T in (fun == BLAS.dot ? RTs : RCs),
+                (sz, inc) in ((10, 1), ((2, 20), -2)),
+                f in (pfun === identity ? (fun, fun_overwrite!) : (fun,))
+
+                Tx <: Const && Ty <: Const && !(Tret <: Const) && continue
+
+                x, ∂x = ntuple(_ -> randn(T, sz), 2)
+                y, ∂y = ntuple(_ -> randn(T, sz), 2)
+                ∂z = randn(T)
+                xcopy, ycopy, ∂xcopy, ∂ycopy = map(copy, (x, y, ∂x, ∂y))
+
+                x_annot =
+                    Tx <: Const ? Const(pfun(xcopy)) : Duplicated(pfun(xcopy), pfun(∂xcopy))
+                y_annot =
+                    Ty <: Const ? Const(pfun(ycopy)) : Duplicated(pfun(ycopy), pfun(∂ycopy))
+                activities = (Const(n), x_annot, Const(inc), y_annot, Const(inc))
+
+                vexp = fun(n, x, inc, y, inc)
+                dret = randn(typeof(vexp))
+
+                dexp = FiniteDifferences.j′vp(
+                    fdm, (x, y) -> fun(n, x, inc, y, inc), dret, x, y
+                )
+                fwd, rev = autodiff_thunk(
+                    ReverseSplitWithPrimal,
+                    Const{typeof(f)},
+                    Tret,
+                    map(typeof, activities)...,
+                )
+                tape, val, shadow_val = fwd(Const(f), activities...)
+                if Tret <: Const
+                    dval, = rev(Const(f), activities..., tape)
+                else
+                    dval, = rev(Const(f), activities..., dret, tape)
+                end
+
+                @test all(isnothing, dval)
+                @test val ≈ vexp
+                @test ∂xcopy ≈
+                    dexp[1] .* !(Tx <: Const || Tret <: Const) .+
+                      ∂x .* ((Tx <: Const) .| (x .== xcopy))
+                @test ∂ycopy ≈
+                    dexp[2] .* !(Ty <: Const || Tret <: Const) .+
+                      ∂y .* ((Ty <: Const) .| (y .== ycopy))
+
+                if pfun === identity && sz == n && inc == 1
+                    @testset "consistency of 2-arg version" begin
+                        xcopy2, ycopy2, ∂xcopy2, ∂ycopy2 = map(copy, (x, y, ∂x, ∂y))
+                        x_annot = if Tx <: Const
+                            Const(pfun(xcopy2))
+                        else
+                            Duplicated(pfun(xcopy2), pfun(∂xcopy2))
+                        end
+                        y_annot = if Ty <: Const
+                            Const(pfun(ycopy2))
+                        else
+                            Duplicated(pfun(ycopy2), pfun(∂ycopy2))
+                        end
+                        activities = (x_annot, y_annot)
+                        fwd, rev = autodiff_thunk(
+                            ReverseSplitWithPrimal,
+                            Const{typeof(f)},
+                            Tret,
+                            map(typeof, activities)...,
+                        )
+                        tape, val2, shadow_val = fwd(Const(f), activities...)
+                        if Tret <: Const
+                            dval2, = rev(Const(f), activities..., tape)
+                        else
+                            dval2, = rev(Const(f), activities..., dret, tape)
+                        end
+                        @test all(isnothing, dval2)
+                        @test val2 == val
+                        @test ∂xcopy2 == ∂xcopy
+                        @test ∂ycopy2 == ∂ycopy
+                    end
+                end
+            end
+        end
+    end
+end

test/runtests.jl

@@ -92,6 +92,7 @@ include("typetree.jl")
         # XXX invalidation does not work on Julia 1.8
         include("ruleinvalidation.jl")
     end
+    include("rules/LinearAlgebra/blas.jl")
 end
 @static if VERSION ≥ v"1.7-" || !Sys.iswindows()
     include("blas.jl")