Merge branch 'JuliaGPU:master' into syrk_strided

Author: evelyne-ringoot

.buildkite/pipeline.yml

@@ -75,7 +75,7 @@ steps:
   - label: "Julia 1.12"
     plugins:
       - JuliaCI/julia#v1:
-          version: "1.12-nightly"
+          version: "1.12"
       - JuliaCI/julia-test#v1:
       - JuliaCI/julia-coverage#v1:
           codecov: true
@@ -111,7 +111,25 @@ steps:
       JULIA_AMDGPU_CORE_MUST_LOAD: "1"
       JULIA_AMDGPU_HIP_MUST_LOAD: "1"
       JULIA_AMDGPU_DISABLE_ARTIFACTS: "1"
-    soft_fail: true
+
+  - label: "Julia 1.11 Enzyme"
+    plugins:
+      - JuliaCI/julia#v1:
+          version: "1.11"
+      - JuliaCI/julia-test#v1:
+          test_args: "enzyme"
+    agents:
+      queue: "juliagpu"
+      rocm: "*"
+      rocmgpu: "*"
+    if: build.message !~ /\[skip tests\]/
+    command: "julia --project -e 'using Pkg; Pkg.update()'"
+    timeout_in_minutes: 30
+    env:
+      JULIA_NUM_THREADS: 4
+      JULIA_AMDGPU_CORE_MUST_LOAD: "1"
+      JULIA_AMDGPU_HIP_MUST_LOAD: "1"
+      JULIA_AMDGPU_DISABLE_ARTIFACTS: "1"
 
   - label: "GPU-less environment"
     plugins:

Project.toml

@@ -1,7 +1,7 @@
 name = "AMDGPU"
 uuid = "21141c5a-9bdb-4563-92ae-f87d6854732e"
 authors = ["Julian P Samaroo <jpsamaroo@jpsamaroo.me>", "Valentin Churavy <v.churavy@gmail.com>", "Anton Smirnov <tonysmn97@gmail.com>"]
-version = "2.1.0"
+version = "2.1.2"
 
 [deps]
 AbstractFFTs = "621f4979-c628-5d54-868e-fcf4e3e8185c"
@@ -52,13 +52,13 @@ EnzymeCore = "0.8"
 ExprTools = "0.1"
 GPUArrays = "11.2.1"
 GPUCompiler = "1"
-GPUToolbox = "0.1.0, 0.2, 0.3"
+GPUToolbox = "0.1.0, 0.2, 0.3, 1"
 KernelAbstractions = "0.9.2"
 LLD_jll = "15, 16, 17, 18, 19"
 LLVM = "9"
 LLVM_jll = "15, 16, 17, 18, 19"
 Preferences = "1"
-PrettyTables = "2"
+PrettyTables = "3"
 ROCmDeviceLibs_jll = "=5.6.1, =6.2.1"
 Random123 = "1.6"
 RandomNumbers = "1.5"

ext/AMDGPUEnzymeCoreExt/AMDGPUEnzymeCoreExt.jl

@@ -14,6 +14,14 @@ function EnzymeCore.compiler_job_from_backend(
     return GPUCompiler.CompilerJob(mi, AMDGPU.compiler_config(AMDGPU.device()))
 end
 
+function EnzymeRules.forward(
+    config, fn::Const{typeof(AMDGPU.hipfunction)}, ::Type{<: Const},
+    f::Const{F}, tt::Const{TT}; kwargs...,
+) where {F, TT}
+    res = fn.val(f.val, tt.val; kwargs...)
+    return res
+end
+
 function EnzymeRules.forward(
     config, fn::Const{typeof(AMDGPU.hipfunction)}, ::Type{<: Duplicated},
     f::Const{F}, tt::Const{TT}; kwargs...,

src/blas/highlevel.jl

@@ -45,6 +45,7 @@ function LinearAlgebra.:(*)(
     dotu(n, x, stride(x, 1), y, stride(y, 1))
 end
 
+LinearAlgebra.norm(x::Diagonal{T, <:StridedROCVector{T}}, p::Real=2) where {T<:Union{Float16, ComplexF16, ROCBLASFloat}} = norm(x.diag, p)
 LinearAlgebra.norm(x::ROCArray{T}) where T <: ROCBLASFloat = nrm2(length(x), x, stride(x, 1))
 LinearAlgebra.BLAS.asum(x::ROCArray{T}) where T <: ROCBLASFloat = asum(length(x), x, stride(x, 1))
 
@@ -163,7 +164,7 @@ if VERSION >= v"1.12-"
     # https://github.com/JuliaLang/LinearAlgebra.jl/blob/4e7c3f40316a956119ac419a97c4b8aad7a17e6c/src/matmul.jl#L490
     for blas_flag in (LinearAlgebra.BlasFlag.SyrkHerkGemm, LinearAlgebra.BlasFlag.SymmHemmGeneric)
         @eval LinearAlgebra.generic_matmatmul_wrapper!(
-            C::StridedROCVecOrMat, tA, tB, A::StridedROCVecOrMat, B::StridedROCVecOrMat,
+            C::StridedROCMatrix, tA::AbstractChar, tB::AbstractChar, A::StridedROCVecOrMat, B::StridedROCVecOrMat,
             alpha::Number, beta::Number, ::$blas_flag) =
             LinearAlgebra.generic_matmatmul!(C, tA, tB, A, B, alpha, beta)
     end
@@ -327,7 +328,7 @@ end
 if VERSION ≥ v"1.12-"
     # Otherwise, dispatches to:
     # https://github.com/JuliaLang/LinearAlgebra.jl/blob/4e7c3f40316a956119ac419a97c4b8aad7a17e6c/src/generic.jl#L2092
-    LinearAlgebra.copytrito!(B::Matrix{T}, A::ROCMatrix{T}, uplo::AbstractChar) where {T <: ROCBLASFloat} = 
+    LinearAlgebra.copytrito!(B::Matrix{T}, A::ROCMatrix{T}, uplo::AbstractChar) where {T <: ROCBLASFloat} =
         invoke(LinearAlgebra.copytrito!, Tuple{AbstractMatrix, AbstractMatrix, AbstractChar}, B, A, uplo)
 end
 

src/blas/util.jl

@@ -1,3 +1,35 @@
+# convert matrix to band storage
+function band(A::AbstractMatrix,kl,ku)
+    m, n = size(A)
+    AB = zeros(eltype(A),kl+ku+1,n)
+    for j = 1:n
+        for i = max(1,j-ku):min(m,j+kl)
+            AB[ku+1-j+i,j] = A[i,j]
+        end
+    end
+    return AB
+end
+
+# convert band storage to general matrix
+function unband(AB::AbstractMatrix,m,kl,ku)
+    bm, n = size(AB)
+    A = zeros(eltype(AB),m,n)
+    for j = 1:n
+        for i = max(1,j-ku):min(m,j+kl)
+            A[i,j] = AB[ku+1-j+i,j]
+        end
+    end
+    return A
+end
+
+# zero out elements not on matrix bands
+function bandex(A::AbstractMatrix,kl,ku)
+    m, n = size(A)
+    AB = band(A,kl,ku)
+    B = unband(AB,m,kl,ku)
+    return B
+end
+
 const ROCBLASReal = Union{Float32, Float64}
 const ROCBLASComplex = Union{ComplexF32, ComplexF64}
 const ROCBLASFloat = Union{ROCBLASReal, ROCBLASComplex}

src/blas/wrappers.jl

@@ -368,10 +368,10 @@ for (fname, elty) in ((:rocblas_stbmv,:Float32),
             x
         end
         function tbmv(
-            uplo::Char, trans::Char, diag::Char,
+            uplo::Char, trans::Char, diag::Char, k::Integer,
             A::ROCMatrix{$elty}, x::ROCVector{$elty},
         )
-            tbmv!(uplo, trans, diag, A, copy(x))
+            tbmv!(uplo, trans, diag, k, A, copy(x))
         end
     end
 end
@@ -496,10 +496,10 @@ for (fname, elty) in ((:rocblas_dsyr,:Float64),
 end
 
 ### her
-for (fname, elty) in ((:rocblas_zher,:ComplexF64),
-                      (:rocblas_cher,:ComplexF32))
+for (fname, elty, relty) in ((:rocblas_zher,:ComplexF64,:Float64),
+                             (:rocblas_cher,:ComplexF32,:Float32))
     @eval begin
-        function her!(uplo::Char, alpha::$elty, x::ROCVector{$elty}, A::ROCMatrix{$elty})
+        function her!(uplo::Char, alpha::$relty, x::ROCVector{$elty}, A::ROCMatrix{$elty})
             m, n = size(A)
             m == n || throw(DimensionMismatch("Matrix A is $m by $n but must be square"))
             length(x) == n || throw(DimensionMismatch("Length of vector must be the same as the matrix dimensions"))
@@ -639,6 +639,7 @@ end
             "- length(B)=$(length(B))\n" *
             "- length(C)=$(length(C))\n"))
     end
+    m, k, n = (-1, -1, -1)
     for (i, (As, Bs, Cs)) in enumerate(zip(A, B, C))
         m, k = size(As, transA == 'N' ? 1 : 2), size(As, transA == 'N' ? 2 : 1)
         n, g = size(Bs, transB == 'N' ? 2 : 1), size(Bs, transB == 'N' ? 1 : 2)
@@ -653,7 +654,7 @@ end
     lda = max(1, stride(A[1], 2))
     ldb = max(1, stride(B[1], 2))
     ldc = max(1, stride(C[1], 2))
-    m, k, n, lda, ldb, ldc
+    return m, k, n, lda, ldb, ldc
 end
 
 ## (GE) general matrix-matrix multiplication batched
@@ -666,15 +667,17 @@ for (fname, elty) in
     @eval begin
         function gemm_batched!(
             transA::Char, transB::Char,
-            alpha::($elty), A::ROCArray{$elty, 3},
-            B::ROCArray{$elty, 3}, beta::($elty), C::ROCArray{$elty, 3},
-        )
+            alpha::($elty), A::TA,
+            B::TB, beta::($elty), C::TC,
+        ) where {TA<:Union{ROCArray{$elty, 3}, Vector{<:ROCMatrix{$elty}}},
+                 TB<:Union{ROCArray{$elty, 3}, Vector{<:ROCMatrix{$elty}}},
+                 TC<:Union{ROCArray{$elty, 3}, Vector{<:ROCMatrix{$elty}}}}
             m, k, n, lda, ldb, ldc = _check_gemm_batched_dims(
                 transA, transB, A, B, C)
 
-            batch_count = size(C, 3)
-            a_broadcast = (size(A, 3) == 1) && (batch_count > 1)
-            b_broadcast = (size(B, 3) == 1) && (batch_count > 1)
+            batch_count = C isa ROCArray ? size(C, 3) : length(C)
+            a_broadcast = A isa ROCArray && (size(A, 3) == 1) && (batch_count > 1)
+            b_broadcast = B isa ROCArray && (size(B, 3) == 1) && (batch_count > 1)
             Ab = a_broadcast ? device_batch(A, batch_count) : device_batch(A)
             Bb = b_broadcast ? device_batch(B, batch_count) : device_batch(B)
             Cb = device_batch(C)
@@ -684,18 +687,18 @@ for (fname, elty) in
                 handle, transA, transB,
                 m, n, k, Ref(alpha), Ab, lda, Bb, ldb, Ref(beta),
                 Cb, ldc, batch_count)
-            C
+            return C
         end
         function gemm_batched(
             transA::Char, transB::Char, alpha::($elty), A::T, B::K,
         ) where {
-            T <: Union{AnyROCArray{$elty, 3}, Vector{ROCMatrix{$elty}}},
-            K <: Union{AnyROCArray{$elty, 3}, Vector{ROCMatrix{$elty}}},
+            T <: Union{AnyROCArray{$elty, 3}, Vector{<:ROCMatrix{$elty}}},
+            K <: Union{AnyROCArray{$elty, 3}, Vector{<:ROCMatrix{$elty}}},
         }
             is_ab_vec = Int(T <: Vector) + Int(K <: Vector)
             (is_ab_vec != 0) && (is_ab_vec != 2) && throw(ArgumentError(
                 "If `A` is a `Vector{ROCMatrix}`, then `B` must be too."))
-            if T isa Vector
+            if T <: Vector
                 C = ROCMatrix{$elty}[similar(B[i], $elty, (
                     size(A[i], transA == 'N' ? 1 : 2),
                     size(B[i], transB == 'N' ? 2 : 1))) for i in 1:length(A)]
@@ -704,13 +707,15 @@ for (fname, elty) in
                 k = size(B, transB == 'N' ? 2 : 1)
                 C = similar(A, $elty, (m, k, max(size(A, 3), size(B, 3))))
             end
-            gemm_batched!(transA, transB, alpha, A, B, zero($elty), C)
+            m, k, n, lda, ldb, ldc = _check_gemm_batched_dims(
+                transA, transB, A, B, C)
+            return gemm_batched!(transA, transB, alpha, A, B, zero($elty), C)
         end
         function gemm_batched(transA::Char, transB::Char, A::T, B::K) where {
-            T <: Union{AnyROCArray{$elty, 3}, Vector{ROCMatrix{$elty}}},
-            K <: Union{AnyROCArray{$elty, 3}, Vector{ROCMatrix{$elty}}},
+            T <: Union{AnyROCArray{$elty, 3}, Vector{<:ROCMatrix{$elty}}},
+            K <: Union{AnyROCArray{$elty, 3}, Vector{<:ROCMatrix{$elty}}},
         }
-            gemm_batched(transA, transB, one($elty), A, B)
+            return gemm_batched(transA, transB, one($elty), A, B)
         end
     end
 end
@@ -863,12 +868,12 @@ for (fname, elty) in ((:rocblas_zhemm,:ComplexF64),
 end
 
 ## herk
-for (fname, elty) in ((:rocblas_zherk,:ComplexF64),
-                      (:rocblas_cherk,:ComplexF32))
+for (fname, elty, relty) in ((:rocblas_zherk,:ComplexF64,:Float64),
+                             (:rocblas_cherk,:ComplexF32,:Float32))
     @eval begin
         function herk!(
-            uplo::Char, trans::Char, alpha::($elty), A::ROCVecOrMat{$elty},
-            beta::($elty), C::ROCMatrix{$elty},
+            uplo::Char, trans::Char, alpha::($relty), A::ROCVecOrMat{$elty},
+            beta::($relty), C::ROCMatrix{$elty},
         )
             mC, n = size(C)
             if mC != n throw(DimensionMismatch("C must be square")) end
@@ -881,12 +886,12 @@ for (fname, elty) in ((:rocblas_zherk,:ComplexF64),
             $(fname)(handle, uplo, trans, n, k, Ref(alpha), A, lda, Ref(beta), C, ldc)
             C
         end
-        function herk(uplo::Char, trans::Char, alpha::($elty), A::ROCVecOrMat{$elty})
+        function herk(uplo::Char, trans::Char, alpha::($relty), A::ROCVecOrMat{$elty})
             n = size(A, trans == 'N' ? 1 : 2)
-            herk!(uplo, trans, alpha, A, zero($elty), similar(A, $elty, (n,n)))
+            herk!(uplo, trans, alpha, A, zero($relty), similar(A, $elty, (n,n)))
         end
         herk(uplo::Char, trans::Char, A::ROCVecOrMat{$elty}) =
-            herk(uplo, trans, one($elty), A)
+            herk(uplo, trans, one($relty), A)
    end
 end
 
@@ -1029,7 +1034,7 @@ for (fname, elty) in
     @eval begin
         function trsm_batched!(
             side::Char, uplo::Char, transa::Char, diag::Char, alpha::($elty),
-            A::Array{ROCMatrix{$elty},1}, B::Array{ROCMatrix{$elty},1},
+            A::Array{<:ROCMatrix{$elty},1}, B::Array{<:ROCMatrix{$elty},1},
         )
             if( length(A) != length(B) )
                 throw(DimensionMismatch(""))
@@ -1051,7 +1056,7 @@ for (fname, elty) in
         end
         function trsm_batched(
             side::Char, uplo::Char, transa::Char, diag::Char, alpha::($elty),
-            A::Array{ROCMatrix{$elty},1}, B::Array{ROCMatrix{$elty},1},
+            A::Array{<:ROCMatrix{$elty},1}, B::Array{<:ROCMatrix{$elty},1},
         )
             trsm_batched!(side, uplo, transa, diag, alpha, A, copy(B) )
         end
@@ -1092,13 +1097,13 @@ for (fname, elty) in ((:rocblas_dgeam,:Float64),
         )
            m,n = size(B)
            if ((transb == 'T' || transb == 'C'))
-               geam!( transa, transb, alpha, A, beta, B, similar(B, $elty, (n,m) ) )
+               return geam!( transa, transb, alpha, A, beta, B, similar(B, $elty, (n,m) ) )
            end
            if (transb == 'N')
-               geam!( transa, transb, alpha, A, beta, B, similar(B, $elty, (m,n) ) )
+               return geam!( transa, transb, alpha, A, beta, B, similar(B, $elty, (m,n) ) )
            end
        end
-       geam( uplo::Char, trans::Char, A::ROCMatrix{$elty}, B::ROCMatrix{$elty}) = geam( uplo, trans, one($elty), A, one($elty), B)
+       geam( transa::Char, transb::Char, A::ROCMatrix{$elty}, B::ROCMatrix{$elty}) = geam( transa, transb, one($elty), A, one($elty), B)
     end
 end
 

src/discovery/discovery.jl

@@ -44,7 +44,7 @@ function _hip_runtime_version()
     VersionNumber(major, minor, patch)
 end
 
-global rel_libdir::String = ""
+global rel_libdir::String = Sys.islinux() ? "" : "bin"
 global libhsaruntime::String = ""
 global lld_path::String = ""
 global lld_artifact::Bool = false

src/hip/device.jl

@@ -168,11 +168,11 @@ function __pretty_data(dev::HIPDevice)
 end
 
 function Base.show(io::IO, mime::MIME{Symbol("text/plain")}, dev::HIPDevice)
-    PrettyTables.pretty_table(io, __pretty_data(dev); header=[
+    PrettyTables.pretty_table(io, __pretty_data(dev); column_labels=[
         "Id", "Name", "GCN arch", "Wavefront", "Memory", "Shared Memory"])
 end
 
 function Base.show(io::IO, mime::MIME{Symbol("text/plain")}, devs::Vector{HIPDevice})
-    PrettyTables.pretty_table(io, vcat(__pretty_data.(devs)...); header=[
+    PrettyTables.pretty_table(io, vcat(__pretty_data.(devs)...); column_labels=[
         "Id", "Name", "GCN arch", "Wavefront", "Memory", "Shared Memory"])
 end

src/utils.jl

@@ -15,7 +15,7 @@ function versioninfo()
         _status(functional(:MIOpen))      "MIOpen"           _ver(:MIOpen, MIOpen.version)       _libpath(libMIOpen_path);
     ]
 
-    PrettyTables.pretty_table(data; header=[
+    PrettyTables.pretty_table(data; column_labels=[
         "Available", "Name", "Version", "Path"],
         alignment=[:c, :l, :l, :l])