DArray: Add diag

docs/src/darray.md

@@ -700,6 +700,7 @@ From `Statistics`:
 
 From `LinearAlgebra`:
 - `norm`
+- `diag`/`tr`
 - `transpose`/`adjoint` (Out-of-place transpose)
 - `*` (Out-of-place Matrix-(Matrix/Vector) multiply)
 - `mul!` (In-place Matrix-Matrix and Matrix-Vector multiply)

src/array/linalg.jl

@@ -5,6 +5,56 @@ function LinearAlgebra.norm2(A::DArray{T,N}) where {T,N}
     return sqrt(sum(map(norm->fetch(norm)::real(T), norms); init=zeroRT))
 end
 
+function _diag_chunk!(Bpart, Brange, Apart, row, col)
+    Ainds = range(LinearIndices(Apart)[row, col]; step=size(Apart, 1)+1, length=length(Brange))
+    view(Bpart, Brange) .= view(Apart, Ainds)
+    return
+end
+
+# `Base`'s fallback indexes with `diagind`, a stepped range of `CartesianIndex`,
+# which `getindex(::DArray, ...)` cannot split across partitions. Copy each
+# partition's stretch of the diagonal instead, so only the partitions that the
+# diagonal crosses are read.
+function LinearAlgebra.diag(A::DMatrix{T}, k::Integer=0) where T
+    m, n = size(A)
+    bs = min(A.partitioning.blocksize...)
+    # Like `diag(::Matrix, k)`, a `k` outside of the matrix gives an empty
+    # diagonal. Compare before converting: `n - k` wraps for an unsigned `k`.
+    -m <= k <= n || return DArray{T}(undef, Blocks(bs), 0)
+    k = Int(k)
+
+    # The k-th diagonal holds the entries (i, i+k)
+    rows = max(1, 1-k):min(m, n-k)
+    B = DArray{T}(undef, Blocks(bs), length(rows))
+    isempty(rows) && return B
+
+    Asd = A.subdomains
+    Bsd = B.subdomains
+    Dagger.spawn_datadeps() do
+        for Aidx in CartesianIndices(A.chunks)
+            Arows, Acols = indexes(Asd[Aidx])
+            # Rows of this partition whose diagonal entry is also in its columns
+            r = intersect(Arows, Acols .- k)
+            isempty(r) && continue
+            Apart = A.chunks[Aidx]
+
+            # Where those entries land in the result
+            p = r .- (first(rows) - 1)
+            for Bidx in fld1(first(p), bs):fld1(last(p), bs)
+                Bpos = only(indexes(Bsd[Bidx]))
+                q = intersect(p, Bpos)
+                Brange = q .- (first(Bpos) - 1)
+                row = first(q) + (first(rows) - 1)
+                Dagger.@spawn _diag_chunk!(InOut(B.chunks[Bidx]), Brange, In(Apart),
+                                           row - (first(Arows) - 1),
+                                           row + k - (first(Acols) - 1))
+            end
+        end
+    end
+
+    return B
+end
+
 # --- BLAS-1 vector operations ---------------------------------------------
 # Efficient, distributed level-1 operations needed by iterative (Krylov)
 # solvers. These run one BLAS-1 kernel per chunk (avoiding the generic

test/array/linalg/core.jl

@@ -77,3 +77,30 @@ end
         @test collect(Dy3) ≈ a .* x .+ b .* y
     end
 end
+
+@testset "diag/tr" begin
+    @testset "$T $(sz) part=$(part.blocksize)" for T in (Float64, ComplexF64),
+            (sz, part) in (((16, 16), Blocks(4, 4)),    # diagonal follows the partition grid
+                           ((16, 16), Blocks(16, 16)),  # single partition
+                           ((16, 16), Blocks(5, 3)),    # diagonal cuts through partitions
+                           ((12, 20), Blocks(4, 5)),    # wide
+                           ((20, 12), Blocks(7, 4)))    # tall
+        A = rand(T, sz...)
+        DA = distribute(A, part)
+        # A `k` outside of the matrix gives an empty diagonal, as for a `Matrix`
+        @testset "k=$k" for k in (-sz[1]-2):(sz[2]+2)
+            Dd = diag(DA, k)
+            @test Dd isa DVector{T}
+            @test collect(Dd) == diag(A, k)
+        end
+        @test collect(diag(DA)) == diag(A)
+        @testset "k::$(typeof(k))" for k in (UInt(1), Int8(-1), Int8(-128), big(2), true, typemax(Int))
+            @test collect(diag(DA, k)) == diag(A, k)
+        end
+        if sz[1] == sz[2]
+            @test tr(DA) ≈ tr(A)
+        else
+            @test_throws DimensionMismatch tr(DA)
+        end
+    end
+end