DArray: CAQR and it's subroutines and helpers

Project.toml

@@ -13,6 +13,7 @@ KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 MacroTools = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
 MemPool = "f9f48841-c794-520a-933b-121f7ba6ed94"
+NextLA = "d37ed344-79c4-486d-9307-6d11355a15a3"
 OnlineStats = "a15396b6-48d5-5d58-9928-6d29437db91e"
 PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a"
 Preferences = "21216c6a-2e73-6563-6e65-726566657250"

src/Dagger.jl

@@ -8,7 +8,7 @@ import MemPool
 import MemPool: DRef, FileRef, poolget, poolset
 
 import Base: collect, reduce
-
+import NextLA
 import LinearAlgebra
 import LinearAlgebra: Adjoint, BLAS, Diagonal, Bidiagonal, Tridiagonal, LAPACK, LowerTriangular, PosDefException, Transpose, UpperTriangular, UnitLowerTriangular, UnitUpperTriangular, diagind, ishermitian, issymmetric
 import Random
@@ -109,6 +109,7 @@ include("array/linalg.jl")
 include("array/mul.jl")
 include("array/cholesky.jl")
 include("array/lu.jl")
+include("array/qr.jl")
 
 import KernelAbstractions, Adapt
 

src/array/alloc.jl

@@ -1,6 +1,7 @@
 import Base: cat
 import Random: MersenneTwister
 export partition
+import LinearAlgebra: UniformScaling
 
 mutable struct AllocateArray{T,N} <: ArrayOp{T,N}
     eltype::Type{T}
@@ -83,17 +84,15 @@ function stage(ctx, a::AllocateArray)
     chunks = map(CartesianIndices(a.domainchunks)) do I
         x = a.domainchunks[I]
         i = LinearIndices(a.domainchunks)[I]
-        args = a.want_index ? (i, size(x)) : (size(x),)
-
         if isnothing(a.procgrid)
             scope = get_compute_scope()
         else
             scope = ExactScope(a.procgrid[CartesianIndex(mod1.(Tuple(I), size(a.procgrid))...)])
         end
         if a.want_index
-            Dagger.@spawn compute_scope=scope allocate_array(a.f, a.eltype, i, args...)
+            Dagger.@spawn compute_scope=scope allocate_array(a.f, a.eltype, i, size(x))
         else
-            Dagger.@spawn compute_scope=scope allocate_array(a.f, a.eltype, args...)
+            Dagger.@spawn compute_scope=scope allocate_array(a.f, a.eltype, size(x))
         end
     end
     return DArray(a.eltype, a.domain, a.domainchunks, chunks, a.partitioning)
@@ -159,6 +158,7 @@ Base.zeros(p::BlocksOrAuto, dims::Dims; assignment::AssignmentType = :arbitrary)
 Base.zeros(::AutoBlocks, eltype::Type, dims::Dims; assignment::AssignmentType = :arbitrary) =
     zeros(auto_blocks(dims), eltype, dims; assignment)
 
+
 function Base.zero(x::DArray{T,N}) where {T,N}
     dims = ntuple(i->x.domain.indexes[i].stop, N)
     sd = first(x.subdomains)
@@ -167,6 +167,39 @@ function Base.zero(x::DArray{T,N}) where {T,N}
     return _to_darray(a)
 end
 
+function _allocate_diag(i,T, _dims, subdomain)
+    sA = zeros(T, _dims)
+    if !isempty(intersect(subdomain.indexes[1], subdomain.indexes[2]))
+        for j in range(1, min(_dims[1], _dims[2]))
+            sA[j,j] = one(T)
+        end
+    end
+    return sA
+end
+
+function DMatrix(p::BlocksOrAuto, s::UniformScaling, dims::Dims, assignment::AssignmentType = :arbitrary)
+    d = ArrayDomain(map(x->1:x, dims))
+    sd = partition(p, d)
+    T = eltype(s)
+    a = AllocateArray(T, (i, T, _dims) -> _allocate_diag(i, T, _dims, sd[i]), true, d, partition(p, d), p, assignment)
+    return _to_darray(a)
+end
+DMatrix(p::BlocksOrAuto, s::UniformScaling, dims::Integer...; assignment::AssignmentType = :arbitrary) = 
+    DMatrix(p, s, dims; assignment)
+DMatrix(::AutoBlocks, s::UniformScaling, dims::Dims; assignment::AssignmentType = :arbitrary) =
+    DMatrix(auto_blocks(dims), s::UniformScaling, dims; assignment)
+
+function DArray{T}(p::BlocksOrAuto, ::UndefInitializer, dims::Dims; assignment::AssignmentType = :arbitrary) where {T}
+    d = ArrayDomain(map(x->1:x, dims))
+    a = AllocateArray(T, AllocateUndef{T}(), false, d, partition(p, d), p, assignment)
+    return _to_darray(a)
+end
+
+DArray{T}(p::BlocksOrAuto, ::UndefInitializer, dims::Integer...; assignment::AssignmentType = :arbitrary) where {T} = 
+    DArray{T}(p, undef, dims; assignment)
+DArray{T}(p::AutoBlocks, ::UndefInitializer, dims::Dims; assignment::AssignmentType = :arbitrary) where {T} = 
+    DArray{T}(auto_blocks(dims), undef, dims; assignment)
+
 function Base.view(A::AbstractArray{T,N}, p::Blocks{N}) where {T,N}
     d = ArrayDomain(Base.index_shape(A))
     dc = partition(p, d)

src/array/copy.jl

@@ -124,7 +124,6 @@ function darray_copyto!(B::DArray{TB,NB}, A::DArray{TA,NA}, Binds=parentindices(
             end
         end
     end
-
     return B
 end
 function copyto_view!(Bpart, Brange, Apart, Arange)

src/array/darray.jl

@@ -65,6 +65,9 @@ ArrayDomain((1:15), (1:80))
 alignfirst(a::ArrayDomain) =
     ArrayDomain(map(r->1:length(r), indexes(a)))
 
+alignfirst(a::CartesianIndices{N}) where N =
+    ArrayDomain(map(r->1:length(r), a.indices))
+
 function size(a::ArrayDomain, dim)
     idxs = indexes(a)
     length(idxs) < dim ? 1 : length(idxs[dim])
@@ -365,7 +368,7 @@ function group_indices(cumlength, idxs::AbstractRange)
 end
 
 _cumsum(x::AbstractArray) = length(x) == 0 ? Int[] : cumsum(x)
-function lookup_parts(A::DArray, ps::AbstractArray, subdmns::DomainBlocks{N}, d::ArrayDomain{N}) where N
+function lookup_parts(A::DArray, ps::AbstractArray, subdmns::DomainBlocks{N}, d::ArrayDomain{N}; slice::Bool=false) where N
     groups = map(group_indices, subdmns.cumlength, indexes(d))
     sz = map(length, groups)
     pieces = Array{Any}(undef, sz)
@@ -379,15 +382,17 @@ function lookup_parts(A::DArray, ps::AbstractArray, subdmns::DomainBlocks{N}, d:
     out_dmn = DomainBlocks(ntuple(x->1,Val(N)), out_cumlength)
     return pieces, out_dmn
 end
-function lookup_parts(A::DArray, ps::AbstractArray, subdmns::DomainBlocks{N}, d::ArrayDomain{S}) where {N,S}
+function lookup_parts(A::DArray, ps::AbstractArray, subdmns::DomainBlocks{N}, d::ArrayDomain{S}; slice::Bool=false) where {N,S}
     if S != 1
         throw(BoundsError(A, d.indexes))
     end
     inds = CartesianIndices(A)[d.indexes...]
     new_d = ntuple(i->first(inds).I[i]:last(inds).I[i], N)
-    return lookup_parts(A, ps, subdmns, ArrayDomain(new_d))
+    return lookup_parts(A, ps, subdmns, ArrayDomain(new_d); slice)
 end
 
+lookup_parts(A::DArray, ps::AbstractArray, subdmns::DomainBlocks{N}, d::CartesianIndices; slice::Bool=false) where {N} = lookup_parts(A, ps, subdmns, ArrayDomain(d.indices); slice)
+
 """
     Base.fetch(c::DArray)
 

src/array/indexing.jl

@@ -1,4 +1,4 @@
-### getindex
+#get index 
 
 const GETINDEX_CACHE = TaskLocalValue{Dict{Tuple,Any}}(()->Dict{Tuple,Any}())
 const GETINDEX_CACHE_SIZE = ScopedValue{Int}(0)
@@ -36,6 +36,7 @@ with_index_caching(f, size::Integer=1) = with(f, GETINDEX_CACHE_SIZE=>size)
     # Return the value
     return part[offset_idx...]
 end
+
 function partition_for(A::DArray, idx::NTuple{N,Int}) where N
     part_idx = zeros(Int, N)
     offset_idx = zeros(Int, N)