DArray: Add proper transpose/adjoint copy/collect

Co-authored-by: Julian P Samaroo <[email protected]>

Author: Rabab53

src/array/matrix.jl

@@ -1,54 +1,22 @@
-struct Transpose{T,N} <: ArrayOp{T,N}
-    f::Function
-    input::ArrayOp
-end
-
-function Transpose(f,x::ArrayOp)
-    @assert 1 <= ndims(x) && ndims(x) <= 2
-    Transpose{eltype(x), 2}(f,x)
-end
-function size(x::Transpose)
-    sz = size(x.input)
-    if length(sz) == 1
-        (1, sz[1])
-    else
-        (sz[2], sz[1])
+# Transpose/Adjoint
+
+function copydiag(f, A::DArray{T, 2}) where T
+    Ac = A.chunks
+    Ac_copy = Matrix{Any}(undef, size(Ac, 2), size(Ac, 1))
+    _copytile(f, Ac) = copy(f(Ac))
+    for i in 1:size(Ac, 1), j in 1:size(Ac, 2)
+        Ac_copy[j, i] = Dagger.@spawn _copytile(f, Ac[i, j])
     end
+    return DArray(T, ArrayDomain(1:size(A,2), 1:size(A,1)), domainchunks(A)', Ac_copy, A.partitioning)
 end
+Base.fetch(A::Adjoint{T, <:DArray{T, 2}}) where T = copydiag(Adjoint, parent(A))
+Base.fetch(A::Transpose{T, <:DArray{T, 2}}) where T = copydiag(Transpose, parent(A))
+Base.copy(A::Adjoint{T, <:DArray{T, 2}}) where T = fetch(A)
+Base.copy(A::Transpose{T, <:DArray{T, 2}}) where T = fetch(A)
+Base.collect(A::Adjoint{T, <:DArray{T, 2}}) where T = collect(copy(A))
+Base.collect(A::Transpose{T, <:DArray{T, 2}}) where T = collect(copy(A))
 
-transpose(x::ArrayOp) = _to_darray(Transpose(transpose, x))
-transpose(x::Union{Chunk, EagerThunk}) = @spawn transpose(x)
-
-adjoint(x::ArrayOp) = _to_darray(Transpose(adjoint, x))
-adjoint(x::Union{Chunk, EagerThunk}) = @spawn adjoint(x)
-
-function adjoint(x::ArrayDomain{2})
-    d = indexes(x)
-    ArrayDomain(d[2], d[1])
-end
-function adjoint(x::ArrayDomain{1})
-    d = indexes(x)
-    ArrayDomain(1:1, d[1])
-end
-
-function adjoint(x::Blocks{2})
-    d = x.blocksize
-    Blocks(d[2], d[1])
-end
-function adjoint(x::Blocks{1})
-    d = x.blocksize
-    Blocks(1, d[1])
-end
-
-function _ctranspose(x::AbstractArray)
-    Any[Dagger.@spawn(adjoint(x[j,i])) for i=1:size(x,2), j=1:size(x,1)]
-end
-
-function stage(ctx::Context, node::Transpose)
-    inp = stage(ctx, node.input)
-    thunks = _ctranspose(chunks(inp))
-    return DArray(eltype(inp), domain(inp)', domainchunks(inp)', thunks, inp.partitioning', inp.concat)
-end
+# Matrix-(Matrix/Vector) multiply
 
 import Base: *, +
 

test/array.jl

@@ -1,5 +1,5 @@
 using LinearAlgebra, SparseArrays, Random, SharedArrays
-import Dagger: chunks, DArray, domainchunks, treereduce_nd
+import Dagger: DArray, chunks, domainchunks, treereduce_nd
 import Distributed: myid, procs
 import Statistics: mean, var, std
 import OnlineStats
@@ -138,7 +138,7 @@ end
 
 @testset "distributing an array" begin
     function test_dist(X)
-        X1 = Distribute(Blocks(10, 20), X)
+        X1 = distribute(X, Blocks(10, 20))
         Xc = fetch(X1)
         @test Xc isa DArray{eltype(X),ndims(X)}
         @test Xc == X
@@ -147,7 +147,7 @@ end
         @test map(x->size(x) == (10, 20), domainchunks(Xc)) |> all
     end
     x = [1 2; 3 4]
-    @test Distribute(Blocks(1,1), x) == x
+    @test distribute(x, Blocks(1,1)) == x
     test_dist(rand(100, 100))
     test_dist(sprand(100, 100, 0.1))
 
@@ -158,7 +158,7 @@ end
 @testset "transpose" begin
     function test_transpose(X)
         x, y = size(X)
-        X1 = Distribute(Blocks(10, 20), X)
+        X1 = distribute(X, Blocks(10, 20))
         @test X1' == X'
         Xc = fetch(X1')
         @test chunks(Xc) |> size == (div(y, 20), div(x,10))
@@ -174,7 +174,7 @@ end
 @testset "matrix-matrix multiply" begin
     function test_mul(X)
         tol = 1e-12
-        X1 = Distribute(Blocks(10, 20), X)
+        X1 = distribute(X, Blocks(10, 20))
         @test_throws DimensionMismatch X1*X1
         X2 = X1'*X1
         X3 = X1*X1'
@@ -188,7 +188,7 @@ end
     test_mul(rand(40, 40))
 
     x = rand(10,10)
-    X = Distribute(Blocks(3,3), x)
+    X = distribute(x, Blocks(3,3))
     y = rand(10)
     @test norm(collect(X*y) - x*y) < 1e-13
 end
@@ -202,24 +202,24 @@ end
 
 @testset "concat" begin
     m = rand(75,75)
-    x = Distribute(Blocks(10,20), m)
-    y = Distribute(Blocks(10,10), m)
+    x = distribute(m, Blocks(10,20))
+    y = distribute(m, Blocks(10,10))
     @test hcat(m,m) == collect(hcat(x,x)) == collect(hcat(x,y))
     @test vcat(m,m) == collect(vcat(x,x))
     @test_throws DimensionMismatch vcat(x,y)
 end
 
 @testset "scale" begin
     x = rand(10,10)
-    X = Distribute(Blocks(3,3), x)
+    X = distribute(x, Blocks(3,3))
     y = rand(10)
 
     @test Diagonal(y)*x == collect(Diagonal(y)*X)
 end
 
 @testset "Getindex" begin
     function test_getindex(x)
-        X = Distribute(Blocks(3,3), x)
+        X = distribute(x, Blocks(3,3))
         @test collect(X[3:8, 2:7]) == x[3:8, 2:7]
         ragged_idx = [1,2,9,7,6,2,4,5]
         @test collect(X[ragged_idx, 2:7]) == x[ragged_idx, 2:7]
@@ -248,7 +248,7 @@ end
 
 
 @testset "cleanup" begin
-    X = Distribute(Blocks(10,10), rand(10,10))
+    X = distribute(rand(10,10), Blocks(10,10))
     @test collect(sin.(X)) == collect(sin.(X))
 end
 
@@ -269,7 +269,7 @@ end
     x=rand(10,10)
     y=copy(x)
     y[3:8, 2:7] .= 1.0
-    X = Distribute(Blocks(3,3), x)
+    X = distribute(x, Blocks(3,3))
     @test collect(setindex(X,1.0, 3:8, 2:7)) == y
     @test collect(X) == x
 end
@@ -292,7 +292,7 @@ end
     @test collect(sort(y)) == x
 
     x = ones(10)
-    y = Distribute(Blocks(3), x)
+    y = distribute(x, Blocks(3))
     @test_broken map(x->length(collect(x)), sort(y).chunks) == [3,3,3,1]
 end