Fixes and tests for banded methods

Author: Katharine Hyatt

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

test/rocarray/blas.jl

@@ -2,7 +2,7 @@
 
 using AMDGPU.rocBLAS
 using AMDGPU.HIP
-import .rocBLAS: rocblas_int
+import .rocBLAS: rocblas_int, bandex, band
 
 m = 20
 n = 35
@@ -231,6 +231,97 @@ end
                 @test B ≈ hB
             end
         end
+        @testset "Banded methods" begin
+            ku = 3
+            kl = 4
+            A  = rand(T, m, n) # generate banded matrix
+            A  = bandex(A, kl, ku)
+            Ab = band(A, kl, ku) # get packed format
+            d_Ab = ROCArray(Ab)
+            x    = rand(T, n)
+            d_x  = ROCArray(x)
+            alpha = rand(T)
+            beta  = rand(T)
+            y     = rand(T, m)
+            d_y   = ROCArray(y)
+            @testset "gbmv!" begin
+                @testset for (op, da1, da2, ha1, ha2)  in (('N',d_x,d_y,x,y), ('T',d_y,d_x,y,x), ('C',d_y,d_x,y,x))
+                    # test y = alpha*A*x + beta*y
+                    rocBLAS.gbmv!(op,m,kl,ku,alpha,d_Ab,da1,beta,da2)
+                    BLAS.gbmv!(op,m,kl,ku,alpha,Ab,ha1,beta,ha2)
+                    @test ha2 ≈ Array(da2)
+                    d_y = ROCArray(y)
+                end
+            end
+            @testset "gbmv" begin
+                d_x = ROCArray(x)
+                d_y = rocBLAS.gbmv('N',m,kl,ku,alpha,d_Ab,d_x)
+                y   = BLAS.gbmv('N',m,kl,ku,alpha,Ab,x)
+                @test y ≈ Array(d_y)
+                # test alpha=1 version without y
+                d_y = rocBLAS.gbmv('N',m,kl,ku,d_Ab,d_x)
+                y   = BLAS.gbmv('N',m,kl,ku,Ab,x)
+                @test y ≈ Array(d_y)
+            end
+            A = rand(T, m, m)
+            A = A + A'
+            nbands = 3
+            @test m >= 1 + nbands
+            A = bandex(A,nbands,nbands)
+            # convert to 'upper' banded storage format
+            AB = band(A, 0, nbands)
+            # construct x
+            x    = rand(T, m)
+            d_AB = ROCArray(AB)
+            d_x  = ROCArray(x)
+            y    = rand(T, m)
+            if T <: Real
+                @testset "sbmv!" begin
+                    d_y = ROCArray(y)
+                    rocBLAS.sbmv!('U',nbands,alpha,d_AB,d_x,beta,d_y)
+                    @test alpha*(A*x) + beta*y ≈ Array(d_y)
+                end
+                @testset "sbmv" begin
+                    d_y = rocBLAS.sbmv('U',nbands,d_AB,d_x)
+                    @test A*x ≈ Array(d_y)
+                end
+            else
+                @testset "hbmv!" begin
+                    d_y = ROCArray(y)
+                    rocBLAS.hbmv!('U',nbands,alpha,d_AB,d_x,beta,d_y)
+                    @test alpha*(A*x) + beta*y ≈ Array(d_y)
+                end
+                @testset "hbmv" begin
+                    d_y = rocBLAS.hbmv('U',nbands,d_AB,d_x)
+                    @test A*x ≈ Array(d_y)
+                end
+            end
+            # generate banded triangular matrix
+            A    = rand(T, m, m)
+            nbands = 3 # restrict to 3 bands
+            @test m >= 1 + nbands
+            A    = bandex(A,0,nbands)
+            AB   = band(A,0,nbands)
+            d_AB = ROCArray(AB)
+            @testset "tbmv!" begin
+                d_y = ROCArray(y)
+                rocBLAS.tbmv!('U','N','N',nbands,d_AB,d_y)
+                @test A*y ≈ Array(d_y)
+            end
+            @testset "tbmv" begin
+                d_y = rocBLAS.tbmv('U','N','N',nbands,d_AB,d_x)
+                @test A*x ≈ Array(d_y)
+            end
+            @testset "tbsv!" begin
+                d_y = copy(d_x)
+                rocBLAS.tbsv!('U','N','N',nbands,d_AB,d_y)
+                @test A\x ≈ Array(d_y)
+            end
+            @testset "tbsv" begin
+                d_y = rocBLAS.tbsv('U','N','N',nbands,d_AB,d_x)
+                @test A\x ≈ Array(d_y)
+            end
+        end
     end
 end