Merge #636

636: Add BLAS.axpby! r=amontoison a=amontoison



Co-authored-by: Alexis Montoison <[email protected]>

src/blas/linalg.jl

@@ -46,6 +46,10 @@ function LinearAlgebra.axpy!(alpha::Number, x::CuArray{T}, y::CuArray{T}) where
     axpy!(length(x), convert(T,alpha), x, 1, y, 1)
 end
 
+function LinearAlgebra.axpby!(alpha::Number, x::CuArray{T}, beta::Number, y::CuArray{T}) where T<:CublasFloat
+    length(x)==length(y) || throw(DimensionMismatch("axpy arguments have lengths $(length(x)) and $(length(y))"))
+    axpby!(length(x), convert(T,alpha), x, 1, convert(T,beta), y, 1)
+end
 
 
 #

src/blas/wrappers.jl

@@ -200,6 +200,33 @@ function axpy!(alpha::Ta,
     y
 end
 
+function axpby!(n::Integer,
+                alpha::T,
+                dx::CuArray{T},
+                incx::Integer,
+                beta::T,
+                dy::CuArray{T},
+                incy::Integer) where T <: CublasFloat
+            scal!(n, beta, dy, incy)
+            axpy!(n, alpha, dx, incx, dy, incy)
+            dy
+end
+
+function axpby!(alpha::Ta,
+                x::CuArray{T},
+                rx::Union{UnitRange{Ti},AbstractRange{Ti}},
+                beta::Tb,
+                y::CuArray{T},
+                ry::Union{UnitRange{Ti},AbstractRange{Ti}}) where {T<:CublasFloat,Ta<:Number,Tb<:Number,Ti<:Integer}
+    length(rx)==length(ry) || throw(DimensionMismatch(""))
+    if minimum(rx) < 1 || maximum(rx) > length(x) || minimum(ry) < 1 || maximum(ry) > length(y)
+        throw(BoundsError())
+    end
+    axpby!(length(rx), convert(T, alpha), pointer(x)+(first(rx)-1)*sizeof(T),
+          step(rx), convert(T, beta), pointer(y)+(first(ry)-1)*sizeof(T), step(ry))
+    y
+end
+
 ## iamax
 # TODO: fix iamax in julia base
 for (fname, elty) in ((:cublasIdamax_v2,:Float64),

test/blas.jl

@@ -32,6 +32,7 @@ k = 13
     @test testf(norm, rand(T, m))
     @test testf(BLAS.asum, rand(T, m))
     @test testf(BLAS.axpy!, Ref(rand()), rand(T, m), rand(T, m))
+    @test testf(BLAS.axpby!, Ref(rand()), rand(T, m), Ref(rand()), rand(T, m))
 
     if T <: Complex
         @test testf(BLAS.dotu, rand(T, m), rand(T, m))