debug

Author: amontoison

src/solver/highlevel.jl

@@ -226,9 +226,9 @@ for (fname, elty) in (
             (mA == nA == mB == nB == mC == nC) || throw(DimensionMismatch("The first two dimensions of A, B and C must match"))
             (nblocksA == nblocksB - 1 == nblocksC) || throw(DimensionMismatch("Inconsistency for the last dimension of A, B and C"))
 
-            lda = max(1, stride(A, 2))
-            ldb = max(1, stride(B, 2))
-            ldc = max(1, stride(C, 2))
+            lda = max(1, stride(A, 3))
+            ldb = max(1, stride(B, 3))
+            ldc = max(1, stride(C, 3))
 
             devinfo = ROCArray{Cint}(undef, 1)
             $fname(rocBLAS.handle(), mB, nblocksB, A, lda, B, ldb, C, ldc, devinfo)
@@ -257,10 +257,10 @@ for (fname, elty) in (
             (mX == mA) || throw(DimensionMismatch("The first dimension of X is inconsistent with first two dimensions of A, B and C"))
             (nblocksA == nblocksB - 1 == nblocksX - 1 == nblocksC) || throw(DimensionMismatch("Inconsistency for the number of blocks in A, B, C and X"))
 
-            lda = max(1, stride(A, 2))
-            ldb = max(1, stride(B, 2))
-            ldc = max(1, stride(C, 2))
-            ldx = max(1, stride(X, 2))
+            lda = max(1, stride(A, 3))
+            ldb = max(1, stride(B, 3))
+            ldc = max(1, stride(C, 3))
+            ldx = max(1, stride(X, 3))
 
             $fname(rocBLAS.handle(), mB, nblocksB, nrhs, A, lda, B, ldb, C, ldc, X, ldx)
             X

test/rocarray/solver.jl

@@ -204,13 +204,15 @@ end
     @testset "elty = $elty" for elty in [Float32, Float64, ComplexF32, ComplexF64]
         nblocks = 5
         n = 16
+        nrhs = 4
         p = n * nblocks
         A = rand(elty, n, n, nblocks-1)
         B = rand(elty, n, n, nblocks)
         C = rand(elty, n, n, nblocks-1)
-        RHS = rand(elty, p, 4)
+        R = rand(elty, n, nblocks, nrhs)
 
         M = zeros(elty, p, p)
+        RHS = zeros(elty, p, nrhs)
         for k in 1:nblocks
             offset = (k-1)*n
             for i = 1:n
@@ -221,14 +223,16 @@ end
                         M[offset+i,offset+n+j] = C[i,j,k]
                     end
                 end
+                for j = 1:nrhs
+                    RHS[offset+i,j] = R[i,k,j]
+                end
             end
         end
-        X = M \ RHS
 
         d_A = ROCArray(A)
         d_B = ROCArray(B)
         d_C = ROCArray(C)
-        d_X = ROCArray(RHS)
+        d_R = ROCArray(R)
         rocSOLVER.geblttrf!(d_A, d_B, d_C)
 
         L = zeros(elty, p, p)
@@ -251,10 +255,21 @@ end
                 end
             end
         end
-        @test L * U ≈ M
+        N = L * U
+        @test N ≈ M
 
-        rocSOLVER.geblttrs!(d_A, d_B, d_C, d_X)
-        @test X ≈ collect(d_X)
+        X = N \ RHS
+        Y = similar(R)
+        for k in 1:nblocks
+            for i = 1:n
+                for j = 1:nrhs
+                    l = (k-1)*n + i
+                    Y[i, k, j] = X[l,j]
+                end
+            end
+        end
+        rocSOLVER.geblttrs!(d_A, d_B, d_C, d_R)
+        @test Y ≈ collect(d_R)
     end
 end