array.jl

export ROCSparseMatrixCSC, ROCSparseMatrixCSR, ROCSparseMatrixBSR, ROCSparseMatrixCOO
export ROCSparseMatrix, AbstractROCSparseMatrix, ROCSparseVector, ROCSparseVecOrMat

abstract type AbstractROCSparseVector{Tv, Ti} <: GPUArrays.AbstractGPUSparseArray{Tv, Ti, 1} end
abstract type AbstractROCSparseMatrix{Tv, Ti} <: GPUArrays.AbstractGPUSparseArray{Tv, Ti, 2} end

mutable struct ROCSparseVector{Tv, Ti} <: AbstractROCSparseVector{Tv, Ti}
    iPtr::ROCVector{Ti}
    nzVal::ROCVector{Tv}
    len::Int
    nnz::Ti

    function ROCSparseVector{Tv, Ti}(
        iPtr::ROCVector{<:Integer}, nzVal::ROCVector, len::Integer,
    ) where {Tv, Ti <: Integer}
        new{Tv, Ti}(iPtr, nzVal, len, length(nzVal))
    end
end

function AMDGPU.unsafe_free!(xs::ROCSparseVector)
    unsafe_free!(nonzeroinds(xs))
    unsafe_free!(nonzeros(xs))
    return
end

mutable struct ROCSparseMatrixCSC{Tv, Ti} <: GPUArrays.AbstractGPUSparseMatrixCSC{Tv, Ti}
    colPtr::ROCVector{Ti}
    rowVal::ROCVector{Ti}
    nzVal::ROCVector{Tv}
    dims::NTuple{2,Int}
    nnz::Ti

    function ROCSparseMatrixCSC{Tv, Ti}(
        colPtr::ROCVector{<:Integer}, rowVal::ROCVector{<:Integer},
        nzVal::ROCVector, dims::NTuple{2,<:Integer},
    ) where {Tv, Ti <: Integer}
        new{Tv, Ti}(colPtr, rowVal, nzVal, dims, length(nzVal))
    end
end
ROCSparseMatrixCSC{Tv, Ti}(x::ROCSparseMatrixCSC{Tv, Ti}) where {Tv, Ti} = x
ROCSparseMatrixCSC(x::ROCSparseMatrixCSC) = x

SparseArrays.rowvals(x::T) where T <: ROCSparseVector = nonzeroinds(x)
SparseArrays.rowvals(x::ROCSparseMatrixCSC) = x.rowVal
SparseArrays.getcolptr(x::ROCSparseMatrixCSC) = x.colPtr

function AMDGPU.unsafe_free!(x::ROCSparseMatrixCSC)
    unsafe_free!(x.colPtr)
    unsafe_free!(rowvals(x))
    unsafe_free!(nonzeros(x))
    return
end

"""
    ROCSparseMatrixCSR{Tv, Ti} <: AbstractROCSparseMatrix{Tv, Ti}

Container to hold sparse matrices in compressed sparse row (CSR) format on the
GPU.

!!! note
    Most ROCSPARSE operations work with CSR formatted matrices, rather
    than CSC.
"""
mutable struct ROCSparseMatrixCSR{Tv, Ti} <: GPUArrays.AbstractGPUSparseMatrixCSR{Tv, Ti}
    rowPtr::ROCVector{Ti}
    colVal::ROCVector{Ti}
    nzVal::ROCVector{Tv}
    dims::NTuple{2,Int}
    nnz::Ti

    function ROCSparseMatrixCSR{Tv, Ti}(
        rowPtr::ROCVector{<:Integer}, colVal::ROCVector{<:Integer},
        nzVal::ROCVector, dims::NTuple{2, Int},
    ) where {Tv, Ti<:Integer}
        new{Tv, Ti}(rowPtr, colVal, nzVal, dims, length(nzVal))
    end
end

ROCSparseMatrixCSR{Tv, Ti}(x::ROCSparseMatrixCSR{Tv, Ti}) where {Tv, Ti} = x
ROCSparseMatrixCSR(x::ROCSparseMatrixCSR) = x

function AMDGPU.unsafe_free!(xs::ROCSparseMatrixCSR)
    unsafe_free!(xs.rowPtr)
    unsafe_free!(xs.colVal)
    unsafe_free!(nonzeros(xs))
    return
end

GPUArrays.sparse_array_type(::Type{<:ROCSparseVector}) = ROCSparseVector
GPUArrays.sparse_array_type(::Type{<:ROCSparseMatrixCSC}) = ROCSparseMatrixCSC
GPUArrays.sparse_array_type(::Type{<:ROCSparseMatrixCSR}) = ROCSparseMatrixCSR
GPUArrays.sparse_array_type(::ROCSparseVector) = ROCSparseVector
GPUArrays.sparse_array_type(::ROCSparseMatrixCSC) = ROCSparseMatrixCSC
GPUArrays.sparse_array_type(::ROCSparseMatrixCSR) = ROCSparseMatrixCSR

GPUArrays.dense_array_type(::Type{<:ROCSparseVector}) = ROCArray
GPUArrays.dense_array_type(::Type{<:ROCSparseMatrixCSC}) = ROCArray
GPUArrays.dense_array_type(::Type{<:ROCSparseMatrixCSR}) = ROCArray
GPUArrays.dense_array_type(::ROCSparseVector) = ROCArray
GPUArrays.dense_array_type(::ROCSparseMatrixCSC) = ROCArray
GPUArrays.dense_array_type(::ROCSparseMatrixCSR) = ROCArray

GPUArrays.csc_type(::ROCSparseMatrixCSR) = ROCSparseMatrixCSC
GPUArrays.csr_type(::ROCSparseMatrixCSC) = ROCSparseMatrixCSR
GPUArrays.coo_type(::Union{ROCSparseMatrixCSR, Transpose{<:Any,<:ROCSparseMatrixCSR}, Adjoint{<:Any,<:ROCSparseMatrixCSR}}) = ROCSparseMatrixCOO
GPUArrays.coo_type(::Union{ROCSparseMatrixCSC, Transpose{<:Any,<:ROCSparseMatrixCSC}, Adjoint{<:Any,<:ROCSparseMatrixCSC}}) = ROCSparseMatrixCOO
GPUArrays.coo_type(::Type{T}) where {T<:Union{ROCSparseMatrixCSR, Transpose{<:Any,<:ROCSparseMatrixCSR}, Adjoint{<:Any,<:ROCSparseMatrixCSR}}} = ROCSparseMatrixCOO
GPUArrays.coo_type(::Type{T}) where {T<:Union{ROCSparseMatrixCSC, Transpose{<:Any,<:ROCSparseMatrixCSC}, Adjoint{<:Any,<:ROCSparseMatrixCSC}}} = ROCSparseMatrixCOO

"""
Container to hold sparse matrices in block compressed sparse row (BSR) format on
the GPU. BSR format is also used in Intel MKL, and is suited to matrices that are
"block" sparse - rare blocks of non-sparse regions.
"""
mutable struct ROCSparseMatrixBSR{Tv, Ti} <: AbstractROCSparseMatrix{Tv, Ti}
    rowPtr::ROCVector{Ti}
    colVal::ROCVector{Ti}
    nzVal::ROCVector{Tv}
    dims::NTuple{2,Int}
    blockDim::Ti
    dir::SparseChar
    nnzb::Ti

    function ROCSparseMatrixBSR{Tv, Ti}(
        rowPtr::ROCVector{<:Integer}, colVal::ROCVector{<:Integer},
        nzVal::ROCVector, dims::NTuple{2,<:Integer},
        blockDim::Integer, dir::SparseChar, nnz::Integer,
    ) where {Tv, Ti<:Integer}
        new{Tv, Ti}(rowPtr, colVal, nzVal, dims, blockDim, dir, nnz)
    end
end

ROCSparseMatrixBSR(A::ROCSparseMatrixBSR) = A

function AMDGPU.unsafe_free!(xs::ROCSparseMatrixBSR)
    unsafe_free!(xs.rowPtr)
    unsafe_free!(xs.colVal)
    unsafe_free!(nonzeros(xs))
    return
end

"""
Container to hold sparse matrices in coordinate (COO) format on the GPU. COO
format is mainly useful to initially construct sparse matrices, afterwards
switch to [`ROCSparseMatrixCSR`](@ref) for more functionality.
"""
mutable struct ROCSparseMatrixCOO{Tv, Ti} <: AbstractROCSparseMatrix{Tv, Ti}
    rowInd::ROCVector{Ti}
    colInd::ROCVector{Ti}
    nzVal::ROCVector{Tv}
    dims::NTuple{2,Int}
    nnz::Ti

    function ROCSparseMatrixCOO{Tv, Ti}(
        rowInd::ROCVector{<:Integer}, colInd::ROCVector{<:Integer},
        nzVal::ROCVector, dims::NTuple{2,Int} = (dimlub(rowInd),dimlub(colInd)),
        nnz::Integer = length(nzVal),
    ) where {Tv, Ti}
        new{Tv, Ti}(rowInd,colInd,nzVal,dims,nnz)
    end
end

ROCSparseMatrixCOO(A::ROCSparseMatrixCOO) = A

"""
Utility union type of [`ROCSparseMatrixCSC`](@ref), [`ROCSparseMatrixCSR`](@ref),
[`ROCSparseMatrixBSR`](@ref), [`ROCSparseMatrixCOO`](@ref).
"""
const ROCSparseMatrix{Tv, Ti} = Union{
    ROCSparseMatrixCSC{Tv, Ti},
    ROCSparseMatrixCSR{Tv, Ti},
    ROCSparseMatrixBSR{Tv, Ti},
    ROCSparseMatrixCOO{Tv, Ti}}

const ROCSparseVecOrMat = Union{ROCSparseVector, ROCSparseMatrix}

# NOTE: we use Cint as default Ti on ROCm instead of Int to provide
# maximum compatiblity to old ROCSparse APIs
# The same pattern was followed for AMDGPU as well
function ROCSparseVector{Tv}(iPtr::ROCVector{<:Integer}, nzVal::ROCVector, len::Integer) where Tv
    ROCSparseVector{Tv, Cint}(convert(ROCVector{Cint}, iPtr), nzVal, len)
end

function ROCSparseMatrixCSC{Tv}(
    colPtr::ROCVector{<:Integer}, rowVal::ROCVector{<:Integer},
    nzVal::ROCVector, dims::NTuple{2,<:Integer},
) where Tv
    ROCSparseMatrixCSC{Tv, Cint}(colPtr, rowVal, nzVal, dims)
end

function ROCSparseMatrixCSR{Tv}(
    rowPtr::ROCVector{<:Integer}, colVal::ROCVector{<:Integer},
    nzVal::ROCVector, dims::NTuple{2,Int},
) where Tv
    ROCSparseMatrixCSR{Tv, Cint}(rowPtr, colVal, nzVal, dims)
end

function ROCSparseMatrixBSR{Tv}(
    rowPtr::ROCVector{<:Integer}, colVal::ROCVector{<:Integer},
    nzVal::ROCVector, dims::NTuple{2,<:Integer},
    blockDim::Integer, dir::SparseChar, nnz::Integer,
) where Tv
    ROCSparseMatrixBSR{Tv, Cint}(rowPtr, colVal, nzVal, dims, blockDim, dir, nnz)
end

function ROCSparseMatrixCOO{Tv}(
    rowInd::ROCVector{<:Integer}, colInd::ROCVector{<:Integer},
    nzVal::ROCVector, dims::NTuple{2,Int} = (dimlub(rowInd), dimlub(colInd)),
    nnz::Integer = length(nzVal),
) where Tv
    ROCSparseMatrixCOO{Tv, Cint}(rowInd,colInd,nzVal,dims,nnz)
end

## convenience constructors
ROCSparseVector(iPtr::DenseROCArray{<:Integer}, nzVal::DenseROCArray{T}, len::Integer) where T =
    ROCSparseVector{T}(iPtr, nzVal, len)

ROCSparseMatrixCSC(
    colPtr::DenseROCArray{<:Integer}, rowVal::DenseROCArray{<:Integer},
    nzVal::DenseROCArray{T}, dims::NTuple{2,Int},
) where {T} =
    ROCSparseMatrixCSC{T}(colPtr, rowVal, nzVal, dims)

ROCSparseMatrixCSR(rowPtr::DenseROCArray, colVal::DenseROCArray, nzVal::DenseROCArray{T}, dims::NTuple{2,Int}) where T =
    ROCSparseMatrixCSR{T}(rowPtr, colVal, nzVal, dims)

ROCSparseMatrixBSR(
    rowPtr::DenseROCArray, colVal::DenseROCArray, nzVal::DenseROCArray{T},
    blockDim, dir, nnz, dims::NTuple{2,Int},
) where T =
    ROCSparseMatrixBSR{T}(rowPtr, colVal, nzVal, dims, blockDim, dir, nnz)

ROCSparseMatrixCOO(
    rowInd::DenseROCArray, colInd::DenseROCArray, nzVal::DenseROCArray{T},
    dims::NTuple{2,Int}, nnz,
) where T =
    ROCSparseMatrixCOO{T}(rowInd, colInd, nzVal, dims, nnz)

Base.similar(Vec::ROCSparseVector) = ROCSparseVector(copy(nonzeroinds(Vec)), similar(nonzeros(Vec)), length(Vec))
Base.similar(Mat::ROCSparseMatrixCSC) = ROCSparseMatrixCSC(copy(Mat.colPtr), copy(rowvals(Mat)), similar(nonzeros(Mat)), size(Mat))
Base.similar(Mat::ROCSparseMatrixCSR) = ROCSparseMatrixCSR(copy(Mat.rowPtr), copy(Mat.colVal), similar(nonzeros(Mat)), size(Mat))
Base.similar(Mat::ROCSparseMatrixBSR) = ROCSparseMatrixBSR(copy(Mat.rowPtr), copy(Mat.colVal), similar(nonzeros(Mat)), Mat.blockDim, Mat.dir, nnz(Mat), size(Mat))
Base.similar(Mat::ROCSparseMatrixCOO) = ROCSparseMatrixCOO(copy(Mat.rowInd), copy(Mat.colInd), similar(nonzeros(Mat)), size(Mat), nnz(Mat))

Base.similar(Vec::ROCSparseVector, T::Type) = ROCSparseVector(copy(nonzeroinds(Vec)), similar(nonzeros(Vec), T), length(Vec))
Base.similar(Mat::ROCSparseMatrixCSC, T::Type) = ROCSparseMatrixCSC(copy(Mat.colPtr), copy(rowvals(Mat)), similar(nonzeros(Mat), T), size(Mat))
Base.similar(Mat::ROCSparseMatrixCSR, T::Type) = ROCSparseMatrixCSR(copy(Mat.rowPtr), copy(Mat.colVal), similar(nonzeros(Mat), T), size(Mat))
Base.similar(Mat::ROCSparseMatrixBSR, T::Type) = ROCSparseMatrixBSR(copy(Mat.rowPtr), copy(Mat.colVal), similar(nonzeros(Mat), T), Mat.blockDim, Mat.dir, nnz(Mat), size(Mat))
Base.similar(Mat::ROCSparseMatrixCOO, T::Type) = ROCSparseMatrixCOO(copy(Mat.rowInd), copy(Mat.colInd), similar(nonzeros(Mat), T), size(Mat), nnz(Mat))

## array interface

Base.length(g::ROCSparseVector) = g.len
Base.size(g::ROCSparseVector) = (g.len,)
Base.ndims(g::ROCSparseVector) = 1

Base.length(g::ROCSparseMatrix) = prod(g.dims)
Base.size(g::ROCSparseMatrix) = g.dims
Base.ndims(g::ROCSparseMatrix) = 2

function Base.size(g::ROCSparseVector, d::Integer)
    if d == 1
        return g.len
    elseif d > 1
        return 1
    else
        throw(ArgumentError("dimension must be ≥ 1, got $d"))
    end
end

function Base.size(g::ROCSparseMatrix, d::Integer)
    if 1 <= d <= 2
        return g.dims[d]
    elseif d > 1
        return 1
    else
        throw(ArgumentError("dimension must be ≥ 1, got $d"))
    end
end

Base.eltype(g::ROCSparseMatrix{T}) where T = T

## sparse array interface

SparseArrays.sparsevec(I::ROCArray{Ti}, V::ROCArray{Tv}, n::Integer) where {Ti,Tv} =
    ROCSparseVector(I, V, n)

function SparseArrays.findnz(S::T) where {T <: AbstractROCSparseMatrix}
    S2 = ROCSparseMatrixCOO(S)
    I = S2.rowInd
    J = S2.colInd
    V = S2.nzVal

    # To make it compatible with the SparseArrays.jl version
    idxs = sortperm(J)
    I = I[idxs]
    J = J[idxs]
    V = V[idxs]

    return (I, J, V)
end

SparseArrays.nnz(g::ROCSparseMatrixBSR) = g.nnzb * g.blockDim * g.blockDim

## indexing

# translations
Base.getindex(A::AbstractROCSparseVector, ::Colon)          = copy(A)
Base.getindex(A::AbstractROCSparseMatrix, ::Colon, ::Colon) = copy(A)
Base.getindex(A::AbstractROCSparseMatrix, i, ::Colon)       = getindex(A, i, 1:size(A, 2))
Base.getindex(A::AbstractROCSparseMatrix, ::Colon, i)       = getindex(A, 1:size(A, 1), i)
Base.getindex(A::AbstractROCSparseMatrix, I::Tuple{Integer,Integer}) = getindex(A, I[1], I[2])

# column slices
function Base.getindex(x::ROCSparseMatrixCSC, ::Colon, j::Integer)
    checkbounds(x, :, j)
    r1 = convert(Int, x.colPtr[j])
    r2 = convert(Int, x.colPtr[j+1]) - 1
    ROCSparseVector(rowvals(x)[r1:r2], nonzeros(x)[r1:r2], size(x, 1))
end

function Base.getindex(x::ROCSparseMatrixCSR, i::Integer, ::Colon)
    checkbounds(x, i, :)
    c1 = convert(Int, x.rowPtr[i])
    c2 = convert(Int, x.rowPtr[i+1]) - 1
    ROCSparseVector(x.colVal[c1:c2], nonzeros(x)[c1:c2], size(x, 2))
end

# row slices
Base.getindex(A::ROCSparseMatrixCSC, i::Integer, ::Colon) = ROCSparseVector(sparse(A[i, 1:end]))  # TODO: optimize
Base.getindex(A::ROCSparseMatrixCSR, ::Colon, j::Integer) = ROCSparseVector(sparse(A[1:end, j]))  # TODO: optimize

function Base.getindex(A::ROCSparseVector{Tv, Ti}, i::Integer) where {Tv, Ti}
    @boundscheck checkbounds(A, i)
    ii = searchsortedfirst(A.iPtr, convert(Ti, i))
    (ii > nnz(A) || A.iPtr[ii] != i) && return zero(Tv)
    A.nzVal[ii]
end

function Base.getindex(A::ROCSparseMatrixCSC{T}, i0::Integer, i1::Integer) where T
    @boundscheck checkbounds(A, i0, i1)
    r1 = Int(A.colPtr[i1])
    r2 = Int(A.colPtr[i1+1]-1)
    (r1 > r2) && return zero(T)
    r1 = searchsortedfirst(rowvals(A), i0, r1, r2, Base.Order.Forward)
    (r1 > r2 || rowvals(A)[r1] != i0) && return zero(T)
    nonzeros(A)[r1]
end

function Base.getindex(A::ROCSparseMatrixCSR{T}, i0::Integer, i1::Integer) where T
    @boundscheck checkbounds(A, i0, i1)
    c1 = Int(A.rowPtr[i0])
    c2 = Int(A.rowPtr[i0+1]-1)
    (c1 > c2) && return zero(T)
    c1 = searchsortedfirst(A.colVal, i1, c1, c2, Base.Order.Forward)
    (c1 > c2 || A.colVal[c1] != i1) && return zero(T)
    nonzeros(A)[c1]
end

function Base.getindex(A::ROCSparseMatrixCOO{T}, i0::Integer, i1::Integer) where T
    @boundscheck checkbounds(A, i0, i1)
    r1 = searchsortedfirst(A.rowInd, i0, Base.Order.Forward)
    (r1 > length(A.rowInd) || A.rowInd[r1] > i0) && return zero(T)
    r2 = searchsortedfirst(A.rowInd, i0+1, Base.Order.Forward)
    c1 = searchsortedfirst(A.colInd, i1, r1, r2, Base.Order.Forward)
    (c1 > r2 || A.colInd[c1] > i1) && return zero(T)
    nonzeros(A)[c1]
end

function Base.getindex(A::ROCSparseMatrixBSR{T}, i0::Integer, i1::Integer) where T
    @boundscheck checkbounds(A, i0, i1)
    i0_block, i0_idx = fldmod1(i0, A.blockDim)
    i1_block, i1_idx = fldmod1(i1, A.blockDim)
    block_idx = (i0_idx - 1) * A.blockDim + i1_idx - 1
    c1 = Int(A.rowPtr[i0_block])
    c2 = Int(A.rowPtr[i0_block+1]-1)
    (c1 > c2) && return zero(T)
    c1 = searchsortedfirst(A.colVal, i1_block, c1, c2, Base.Order.Forward)
    (c1 > c2 || A.colVal[c1] != i1_block) && return zero(T)
    nonzeros(A)[c1+block_idx]
end

## interop with sparse CPU arrays

# cpu to gpu
# NOTE: we eagerly convert the indices to Cint here to avoid additional conversion later on
ROCSparseVector{T}(Vec::SparseVector) where {T} =
    ROCSparseVector(ROCVector{Cint}(Vec.nzind), ROCVector{T}(Vec.nzval), length(Vec))
ROCSparseVector{T}(Mat::SparseMatrixCSC) where {T} =
    size(Mat,2) == 1 ?
        ROCSparseVector(ROCVector{Cint}(Mat.rowval), ROCVector{T}(Mat.nzval), size(Mat)[1]) :
        throw(ArgumentError("The input argument must have a single column"))
ROCSparseMatrixCSC{T}(Vec::SparseVector) where {T} =
    ROCSparseMatrixCSC{T}(
        ROCVector{Cint}([1]), ROCVector{Cint}(Vec.nzind),
        ROCVector{T}(Vec.nzval), size(Vec))
ROCSparseMatrixCSC{T}(Mat::SparseMatrixCSC) where {T} =
    ROCSparseMatrixCSC{T}(
        ROCVector{Cint}(Mat.colptr), ROCVector{Cint}(Mat.rowval),
        ROCVector{T}(Mat.nzval), size(Mat))
ROCSparseMatrixCSR{T}(Mat::Transpose{Tv, <:SparseMatrixCSC}) where {T, Tv} =
    ROCSparseMatrixCSR{T}(
        ROCVector{Cint}(parent(Mat).colptr), ROCVector{Cint}(parent(Mat).rowval),
        ROCVector{T}(parent(Mat).nzval), size(Mat))
ROCSparseMatrixCSR{T}(Mat::Adjoint{Tv, <:SparseMatrixCSC}) where {T, Tv} =
    ROCSparseMatrixCSR{T}(
        ROCVector{Cint}(parent(Mat).colptr), ROCVector{Cint}(parent(Mat).rowval),
        ROCVector{T}(conj.(parent(Mat).nzval)), size(Mat))
ROCSparseMatrixCSC{T}(Mat::Union{Transpose{Tv, <:SparseMatrixCSC}, Adjoint{Tv, <:SparseMatrixCSC}}) where {T, Tv} = ROCSparseMatrixCSC(ROCSparseMatrixCSR{T}(Mat))
ROCSparseMatrixCSR{T}(Mat::SparseMatrixCSC) where {T} = ROCSparseMatrixCSR(ROCSparseMatrixCSC{T}(Mat))
ROCSparseMatrixBSR{T}(Mat::SparseMatrixCSC, blockdim) where {T} = ROCSparseMatrixBSR(ROCSparseMatrixCSR{T}(Mat), blockdim)
ROCSparseMatrixCOO{T}(Mat::SparseMatrixCSC) where {T} = ROCSparseMatrixCOO(ROCSparseMatrixCSR{T}(Mat))

ROCSparseMatrixCOO{T}(Mat::Transpose{Tv, <:SparseMatrixCSC}) where {T, Tv} = ROCSparseMatrixCOO{T}(ROCSparseMatrixCSR{T}(Mat))
ROCSparseMatrixCOO{T}(Mat::Adjoint{Tv, <:SparseMatrixCSC}) where {T, Tv} = ROCSparseMatrixCOO{T}(ROCSparseMatrixCSR{T}(Mat))

# untyped variants
ROCSparseVector(x::AbstractSparseArray{T}) where {T} = ROCSparseVector{T}(x)
ROCSparseMatrixCSC(x::AbstractSparseArray{T}) where {T} = ROCSparseMatrixCSC{T}(x)
ROCSparseMatrixCSR(x::AbstractSparseArray{T}) where {T} = ROCSparseMatrixCSR{T}(x)
ROCSparseMatrixBSR(x::AbstractSparseArray{T}, blockdim) where {T} = ROCSparseMatrixBSR{T}(x, blockdim)
ROCSparseMatrixCOO(x::AbstractSparseArray{T}) where {T} = ROCSparseMatrixCOO{T}(x)

# adjoint / transpose
ROCSparseMatrixCSR(x::Transpose{T}) where {T} = ROCSparseMatrixCSR{T}(x)
ROCSparseMatrixCSR(x::Adjoint{T}) where {T} = ROCSparseMatrixCSR{T}(x)
ROCSparseMatrixCSC(x::Transpose{T}) where {T} = ROCSparseMatrixCSC{T}(x)
ROCSparseMatrixCSC(x::Adjoint{T}) where {T} = ROCSparseMatrixCSC{T}(x)
ROCSparseMatrixCOO(x::Transpose{T}) where {T} = ROCSparseMatrixCOO{T}(x)
ROCSparseMatrixCOO(x::Adjoint{T}) where {T} = ROCSparseMatrixCOO{T}(x)

ROCSparseMatrixCSR(x::Transpose{T,<:Union{ROCSparseMatrixCSC, ROCSparseMatrixCSR, ROCSparseMatrixCOO}}) where {T} = ROCSparseMatrixCSR(_sptranspose(parent(x)))
ROCSparseMatrixCSC(x::Transpose{T,<:Union{ROCSparseMatrixCSC, ROCSparseMatrixCSR, ROCSparseMatrixCOO}}) where {T} = ROCSparseMatrixCSC(_sptranspose(parent(x)))
ROCSparseMatrixCOO(x::Transpose{T,<:Union{ROCSparseMatrixCSC, ROCSparseMatrixCSR, ROCSparseMatrixCOO}}) where {T} = ROCSparseMatrixCOO(_sptranspose(parent(x)))
ROCSparseMatrixCSR(x::Adjoint{T,<:Union{ROCSparseMatrixCSC, ROCSparseMatrixCSR, ROCSparseMatrixCOO}}) where {T} = ROCSparseMatrixCSR(_spadjoint(parent(x)))
ROCSparseMatrixCSC(x::Adjoint{T,<:Union{ROCSparseMatrixCSC, ROCSparseMatrixCSR, ROCSparseMatrixCOO}}) where {T} = ROCSparseMatrixCSC(_spadjoint(parent(x)))
ROCSparseMatrixCOO(x::Adjoint{T,<:Union{ROCSparseMatrixCSC, ROCSparseMatrixCSR, ROCSparseMatrixCOO}}) where {T} = ROCSparseMatrixCOO(_spadjoint(parent(x)))

# gpu to cpu
SparseVector(x::ROCSparseVector) = SparseVector(length(x), Array(SparseArrays.nonzeroinds(x)), Array(SparseArrays.nonzeros(x)))
SparseMatrixCSC(x::ROCSparseMatrixCSC) = SparseMatrixCSC(size(x)..., Array(x.colPtr), Array(SparseArrays.rowvals(x)), Array(SparseArrays.nonzeros(x)))
SparseMatrixCSC(x::ROCSparseMatrixCSR) = SparseMatrixCSC(ROCSparseMatrixCSC(x))  # no direct conversion
SparseMatrixCSC(x::ROCSparseMatrixBSR) = SparseMatrixCSC(ROCSparseMatrixCSR(x))  # no direct conversion
SparseMatrixCSC(x::ROCSparseMatrixCOO) = SparseMatrixCSC(ROCSparseMatrixCSR(x))  # no direct conversion

Adapt.adapt_storage(::Type{ROCArray}, xs::SparseVector) = ROCSparseVector(xs)
Adapt.adapt_storage(::Type{ROCArray}, xs::SparseMatrixCSC) = ROCSparseMatrixCSC(xs)
Adapt.adapt_storage(::Type{ROCArray{T}}, xs::SparseVector) where {T} = ROCSparseVector{T}(xs)
Adapt.adapt_storage(::Type{ROCArray{T}}, xs::SparseMatrixCSC) where {T} = ROCSparseMatrixCSC{T}(xs)

Adapt.adapt_storage(::AMDGPU.Float32Adaptor, xs::AbstractSparseArray) =
  adapt(ROCArray, xs)
Adapt.adapt_storage(::AMDGPU.Float32Adaptor, xs::AbstractSparseArray{<:AbstractFloat}) =
  adapt(ROCArray{Float32}, xs)

Adapt.adapt_storage(::Type{Array}, xs::ROCSparseVector) = SparseVector(xs)
Adapt.adapt_storage(::Type{Array}, xs::ROCSparseMatrixCSC) = SparseMatrixCSC(xs)

## copying between sparse GPU arrays

function Base.copyto!(dst::ROCSparseMatrixCSR, src::ROCSparseMatrixCSR)
    if size(dst) != size(src)
        throw(ArgumentError("Inconsistent Sparse Matrix size"))
    end
    copyto!(dst.rowPtr, src.rowPtr)
    copyto!(dst.colVal, src.colVal)
    copyto!(nonzeros(dst), nonzeros(src))
    dst.nnz = src.nnz
    dst
end

function Base.copyto!(dst::ROCSparseMatrixBSR, src::ROCSparseMatrixBSR)
    if size(dst) != size(src)
        throw(ArgumentError("Inconsistent Sparse Matrix size"))
    end
    copyto!(dst.rowPtr, src.rowPtr)
    copyto!(dst.colVal, src.colVal)
    copyto!(nonzeros(dst), nonzeros(src))
    dst.dir = src.dir
    dst.nnzb = src.nnzb
    dst
end

function Base.copyto!(dst::ROCSparseMatrixCOO, src::ROCSparseMatrixCOO)
    if size(dst) != size(src)
        throw(ArgumentError("Inconsistent Sparse Matrix size"))
    end
    copyto!(dst.rowInd, src.rowInd)
    copyto!(dst.colInd, src.colInd)
    copyto!(nonzeros(dst), nonzeros(src))
    dst.nnz = src.nnz
    dst
end

Base.copy(Vec::ROCSparseVector) = copyto!(similar(Vec), Vec)
Base.copy(Mat::ROCSparseMatrixCSC) = copyto!(similar(Mat), Mat)
Base.copy(Mat::ROCSparseMatrixCSR) = copyto!(similar(Mat), Mat)
Base.copy(Mat::ROCSparseMatrixBSR) = copyto!(similar(Mat), Mat)
Base.copy(Mat::ROCSparseMatrixCOO) = copyto!(similar(Mat), Mat)

# input/output

for (gpu, cpu) in [:ROCSparseVector => :SparseVector]
    @eval function Base.show(io::IO, ::MIME"text/plain", x::$gpu)
        xnnz = length(nonzeros(x))
        print(io, length(x), "-element ", typeof(x), " with ", xnnz,
            " stored ", xnnz == 1 ? "entry" : "entries")
        if xnnz != 0
            println(io, ":")
            show(IOContext(io, :typeinfo => eltype(x)), $cpu(x))
        end
    end
end

for (gpu, cpu) in [:ROCSparseMatrixCSC => :SparseMatrixCSC,
                   :ROCSparseMatrixCSR => :SparseMatrixCSC,
                   :ROCSparseMatrixBSR => :SparseMatrixCSC,
                   :ROCSparseMatrixCOO => :SparseMatrixCSC]
    @eval Base.show(io::IOContext, x::$gpu) = show(io, $cpu(x))

    @eval function Base.show(io::IO, mime::MIME"text/plain", S::$gpu)
        xnnz = nnz(S)
        m, n = size(S)
        print(io, m, "×", n, " ", typeof(S), " with ", xnnz, " stored ",
                  xnnz == 1 ? "entry" : "entries")
        if !(m == 0 || n == 0)
            println(io, ":")
            io = IOContext(io, :typeinfo => eltype(S))
            if ndims(S) == 1
                show(io, $cpu(S))
            else
                # so that we get the nice Braille pattern
                Base.print_array(io, $cpu(S))
            end
        end
    end
end

# interop with device arrays

Adapt.adapt_structure(to::AMDGPU.Runtime.Adaptor, x::ROCSparseVector) =
    GPUArrays.GPUSparseDeviceVector(adapt(to, x.iPtr), adapt(to, x.nzVal), length(x), x.nnz)

Adapt.adapt_structure(to::AMDGPU.Runtime.Adaptor, x::ROCSparseMatrixCSR) =
    GPUArrays.GPUSparseDeviceMatrixCSR(
        adapt(to, x.rowPtr), adapt(to, x.colVal), adapt(to, x.nzVal), size(x), x.nnz)

Adapt.adapt_structure(to::AMDGPU.Runtime.Adaptor, x::ROCSparseMatrixCSC) =
    GPUArrays.GPUSparseDeviceMatrixCSC(
        adapt(to, x.colPtr), adapt(to, x.rowVal), adapt(to, x.nzVal), size(x), x.nnz)

Adapt.adapt_structure(to::AMDGPU.Runtime.Adaptor, x::ROCSparseMatrixBSR) =
    GPUArrays.GPUSparseDeviceMatrixBSR(
        adapt(to, x.rowPtr), adapt(to, x.colVal), adapt(to, x.nzVal),
        size(x), x.blockDim, x.dir, x.nnzb)

Adapt.adapt_structure(to::AMDGPU.Runtime.Adaptor, x::ROCSparseMatrixCOO) =
    GPUArrays.GPUSparseDeviceMatrixCOO(
        adapt(to, x.rowInd), adapt(to, x.colInd), adapt(to, x.nzVal), size(x), x.nnz)

# device array ctors

GPUArrays.GPUSparseDeviceVector(
    iPtr::ROCDeviceVector{Ti, A}, nzVal::ROCDeviceVector{Tv, A}, len::Int, nnz::Ti,
) where {Ti, Tv, A} = GPUArrays.GPUSparseDeviceVector{
    Tv, Ti, ROCDeviceVector{Ti, A}, ROCDeviceVector{Tv, A}, A,
}(iPtr, nzVal, len, nnz)

GPUArrays.GPUSparseDeviceMatrixCSR(
    rowPtr::ROCDeviceVector{Ti, A}, colVal::ROCDeviceVector{Ti, A}, nzVal::ROCDeviceVector{Tv, A},
    dims::NTuple{2, Int}, nnz::Ti,
) where {Ti, Tv, A} = GPUArrays.GPUSparseDeviceMatrixCSR{
    Tv, Ti, ROCDeviceVector{Ti, A}, ROCDeviceVector{Tv, A}, A,
}(rowPtr, colVal, nzVal, dims, nnz)

GPUArrays.GPUSparseDeviceMatrixCSC(
    colPtr::ROCDeviceVector{Ti, A}, rowVal::ROCDeviceVector{Ti, A}, nzVal::ROCDeviceVector{Tv, A},
    dims::NTuple{2, Int}, nnz::Ti,
) where {Ti, Tv, A} = GPUArrays.GPUSparseDeviceMatrixCSC{
    Tv, Ti, ROCDeviceVector{Ti, A}, ROCDeviceVector{Tv, A}, A,
}(colPtr, rowVal, nzVal, dims, nnz)

GPUArrays.GPUSparseDeviceMatrixBSR(
    rowPtr::ROCDeviceVector{Ti, A}, colVal::ROCDeviceVector{Ti, A}, nzVal::ROCDeviceVector{Tv, A},
    dims::NTuple{2, Int}, blockDim::Ti, dir::Char, nnz::Ti,
) where {Ti, Tv, A} = GPUArrays.GPUSparseDeviceMatrixBSR{
    Tv, Ti, ROCDeviceVector{Ti, A}, ROCDeviceVector{Tv, A}, A,
}(rowPtr, colVal, nzVal, dims, blockDim, dir, nnz)

GPUArrays.GPUSparseDeviceMatrixCOO(
    rowInd::ROCDeviceVector{Ti, A}, colInd::ROCDeviceVector{Ti, A}, nzVal::ROCDeviceVector{Tv, A},
    dims::NTuple{2, Int}, nnz::Ti,
) where {Ti, Tv, A} = GPUArrays.GPUSparseDeviceMatrixCOO{
    Tv, Ti, ROCDeviceVector{Ti, A}, ROCDeviceVector{Tv, A}, A,
}(rowInd, colInd, nzVal, dims, nnz)