diff --git a/lib/cublas/linalg.jl b/lib/cublas/linalg.jl index c8b205a2a5..a959fad268 100644 --- a/lib/cublas/linalg.jl +++ b/lib/cublas/linalg.jl @@ -37,6 +37,16 @@ end LinearAlgebra.norm(x::DenseCuArray{<:CublasFloat}) = nrm2(x) LinearAlgebra.BLAS.asum(x::StridedCuArray{<:CublasFloat}) = asum(length(x), x) +function LinearAlgebra.norm(x::DenseCuArray{<:CublasFloat}, p::Integer) + if p==1 + return CUBLAS.asum(length(x),x) + elseif p==2 + return LinearAlgebra.norm(x) + else + return LinearAlgebra.tr(LinearAlgebra.Diagonal(abs.(x))^p)^(1/p) + end +end + function LinearAlgebra.axpy!(alpha::Number, x::StridedCuArray{T}, y::StridedCuArray{T}) where T<:CublasFloat length(x)==length(y) || throw(DimensionMismatch("axpy arguments have lengths $(length(x)) and $(length(y))")) axpy!(length(x), alpha, x, y) diff --git a/test/cublas.jl b/test/cublas.jl index 4319abb6bc..3e00e64d80 100644 --- a/test/cublas.jl +++ b/test/cublas.jl @@ -1,5 +1,6 @@ using CUDA.CUBLAS using CUDA.CUBLAS: band, bandex +using GPUArrays using LinearAlgebra @@ -84,6 +85,14 @@ end mul!(y, f(A), x, Ts(1), Ts(2)) @test Array(dy) ≈ y end + + @testset "norm" begin + x, y, z = CUDA.rand(elty, 1), CUDA.rand(elty, 2), CUDA.rand(elty, m) + @disallowscalar @test typeof(norm(x,1)) <: Real + @disallowscalar @test typeof(norm(y,2)) <: Real + @disallowscalar @test typeof(norm(z,m)) <: Real + end + @testset "banded methods" begin # bands ku = 2