Add findnz & fix broadcasting (#704)

Project.toml

@@ -1,7 +1,7 @@
 name = "AMDGPU"
 uuid = "21141c5a-9bdb-4563-92ae-f87d6854732e"
 authors = ["Julian P Samaroo <[email protected]>", "Valentin Churavy <[email protected]>", "Anton Smirnov <[email protected]>"]
-version = "1.1.2"
+version = "1.1.3"
 
 [deps]
 AbstractFFTs = "621f4979-c628-5d54-868e-fcf4e3e8185c"

src/broadcast.jl

@@ -9,7 +9,18 @@ BroadcastStyle(::Type{<:ROCArray{T, N, B}}) where {T, N, B} =
 BroadcastStyle(W::Type{<:AnyROCArray{T, N}}) where {T, N} =
     ROCArrayStyle{N, buftype(Adapt.unwrap_type(W))}()
 
-# TODO handle broadcast of different buffer types (use unified memory).
+# TODO use unified buffer once we support it.
+# Broadcast of two different buffers - choose `HIPBuffer`.
+BroadcastStyle(
+    ::ROCArrayStyle{N1, B1},
+    ::ROCArrayStyle{N2, B2},
+) where {N1,N2,B1,B2} = ROCArrayStyle{max(N1,N2), Mem.HIPBuffer}()
+
+# Different N, same buffer type.
+BroadcastStyle(
+    ::ROCArrayStyle{N1, B},
+    ::ROCArrayStyle{N2, B},
+) where {N1,N2,B} = ROCArrayStyle{max(N1,N2), B}()
 
 # Allocation of output arrays.
 function Base.similar(

src/sparse/array.jl

@@ -262,12 +262,31 @@ 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::AbstractROCSparseArray) = g.nnz
 SparseArrays.nonzeros(g::AbstractROCSparseArray) = g.nzVal
 
 SparseArrays.nonzeroinds(g::AbstractROCSparseVector) = g.iPtr
 
 SparseArrays.rowvals(g::ROCSparseMatrixCSC) = g.rowVal
+SparseArrays.getcolptr(g::ROCSparseMatrixCSC) = g.colPtr
 
 LinearAlgebra.issymmetric(M::Union{ROCSparseMatrixCSC,ROCSparseMatrixCSR}) = false
 LinearAlgebra.ishermitian(M::Union{ROCSparseMatrixCSC,ROCSparseMatrixCSR}) = false

test/rocarray/base.jl

@@ -145,6 +145,14 @@ end
         AMDGPU.unsafe_free!(xd2)
         @test_throws ArgumentError pointer(xd2)
     end
+
+    @testset "Broadcasting different buffer types" begin
+        x = rand(Float32, 4, 16, 16)
+        xd = unsafe_wrap(ROCArray, pointer(x), size(x))
+        y = AMDGPU.zeros(Float32, 3, 16, 16)
+        y .= @view(xd[1:3, :, :])
+        @test Array(y) ≈ @view(x[1:3, :, :])
+    end
 end
 
 @testset "unsafe_free" begin

test/rocsparse/interfaces.jl

@@ -203,3 +203,26 @@
         @test (dD - dA) isa typ
     end
 end
+
+@testset "SparseArrays.jl" begin
+    @testset "findnz" begin
+        n = 35
+        A = sprand(n, n, 0.2)
+        d_A = ROCSparseMatrixCSC(A)
+        @test Array(SparseArrays.getcolptr(d_A)) == SparseArrays.getcolptr(A)
+
+        i, j, v = findnz(A)
+        d_i, d_j, d_v = findnz(d_A)
+        @test Array(d_i) == i && Array(d_j) == j && Array(d_v) == v
+
+        i = unique(sort(rand(1:n, 10)))
+        vals = rand(length(i))
+        d_i = ROCArray(i)
+        d_vals = ROCArray(vals)
+        v = sparsevec(i, vals, n)
+        d_v = sparsevec(d_i, d_vals, n)
+        @test Array(d_v.iPtr) == v.nzind
+        @test Array(d_v.nzVal) == v.nzval
+        @test d_v.len == v.n
+    end
+end