unify device_batch for CUDA/ROCm

Author: jindavid

src/backbone_cuda.jl

@@ -1,6 +1,11 @@
 using CUDA
 using CUDA.CUSOLVER, CUDA.CUBLAS
 
+@inline function device_batch(batch::Vector{<:CuArray{T}}) where {T}
+    ptrs = pointer.(batch)
+    return CuArray(ptrs)
+end
+
 for (Xpotrf_buffer, Xpotrf, Xtrsm, Xgemm, T) in (
     (:cusolverDnSpotrf_bufferSize, :cusolverDnSpotrf, :cublasStrsm_v2, :cublasSgemm_v2, :Float32),
     (:cusolverDnDpotrf_bufferSize, :cusolverDnDpotrf, :cublasDtrsm_v2, :cublasDgemm_v2, :Float64),

src/backbone_rocm.jl

@@ -1,6 +1,11 @@
 using AMDGPU
 using AMDGPU.rocSOLVER, AMDGPU.rocBLAS
 
+function device_batch(batch::Array{T}) where T <: ROCArray
+    E = eltype(T)
+    ROCArray([convert(Ptr{E}, arr.buf[]) for arr in batch])
+end
+
 for (Xpotrf, Xtrsm, Xgemm, T) in (
     (:rocsolver_spotrf, :rocblas_strsm, :rocblas_sgemm, :Float32),
     (:rocsolver_dpotrf, :rocblas_dtrsm, :rocblas_dgemm, :Float64),
@@ -151,7 +156,7 @@ for (XpotrfBatched, XtrsmBatched, XgemmBatched, T) in (
         function update_boundary!(M_ptrs_1::ROCVector{<:Ptr{$T}}, M_ptrs_2::ROCVector{<:Ptr{$T}}, d_ptrs::ROCVector{<:Ptr{$T}}, P, n, m)
 
             dh = rocBLAS.handle()
-            
+
             rocBLAS.$XgemmBatched(
                 dh, rocBLAS.rocblas_operation_transpose, rocBLAS.rocblas_operation_none,
                 n, 1, n, -one($T),

src/gpu_seq.jl

@@ -28,16 +28,16 @@ end
 
 function create_matrix_list(N::Int, n1::Int, n2::Int, ::Type{T}, ::Type{M}) where {T, M}
 
-    M_vec = M{T, 2}(zeros(N*n1*n2, 1)) #TODO
-    M_tensor = unsafe_wrap(M{T, 3}, pointer(M_vec), (n1, n2, N); own=false) #TODO
+    M_vec = M{T, 2}(zeros(N*n1*n2, 1))
+    M_tensor = unsafe_wrap(M{T, 3}, pointer(M_vec), (n1, n2, N); own=false)
     M_list = Vector{M{T, 2}}(undef, N);
     ptr = pointer(M_tensor)
 
     for i in 1:N
         M_list[i] = unsafe_wrap(M{T, 2}, ptr + n1*n2*(i-1)*sizeof(T), (n1, n2); own=false)
     end
 
-    M_ptrs = CUBLAS.unsafe_batch(M_list) #TODO
+    M_ptrs = device_batch(M_list)
 
     return M_vec, M_tensor, M_list, M_ptrs
 end
@@ -66,7 +66,7 @@ function factorize!(data::BlockTriDiagData_seq)
     A_ptrs = data.A_ptrs
     B_ptrs = data.B_ptrs
 
-    @allowscalar cholesky_factorize!(A_ptrs, B_ptrs, N, n) #TODO
+    @allowscalar cholesky_factorize!(A_ptrs, B_ptrs, N, n) #TODO check if works for both CUDA and ROCm
 
 end
 
@@ -79,6 +79,6 @@ function solve!(data::BlockTriDiagData_seq)
     B_ptrs = data.B_ptrs
     d_ptrs = data.d_ptrs
 
-    @allowscalar cholesky_solve!(A_ptrs, B_ptrs, d_ptrs, N, n, 1) #TODO
+    @allowscalar cholesky_solve!(A_ptrs, B_ptrs, d_ptrs, N, n, 1) #TODO check if works for both CUDA and ROCm
 
 end