WIP: snapshot current local changes

Author: fda-tome

src/array/alloc.jl

@@ -1,7 +1,6 @@
 import Base: cat
 import Random: MersenneTwister, rand!, randn!
 export partition
-import LinearAlgebra: UniformScaling
 
 mutable struct AllocateArray{T,N} <: ArrayOp{T,N}
     eltype::Type{T}
@@ -177,7 +176,6 @@ Base.zeros(p::BlocksOrAuto, dims::Dims; assignment::AssignmentType = :arbitrary)
 Base.zeros(::AutoBlocks, T::Type, dims::Dims; assignment::AssignmentType = :arbitrary) =
     zeros(auto_blocks(dims), T, 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)
@@ -198,60 +196,6 @@ function LinearAlgebra._zeros(::Type{T}, B::DMatrix, n::Integer) where T
     return zeros(auto_blocks(sz), T, sz)
 end
 
-function _allocate_diag(i,T, _dims, domain, p)
-    sA = zeros(T, _dims)
-    dom_idx = indexes(domain)
-    chunk_counts = ntuple(j -> cld(length(dom_idx[j]), p.blocksize[j]), length(dom_idx))
-    subinds_idx = CartesianIndices(ntuple(j -> Base.OneTo(chunk_counts[j]), length(chunk_counts)))[i]
-    subinds = Tuple(subinds_idx)
-    sd = map(enumerate(dom_idx)) do (idx, range)
-        bsz = div(last(range), p.blocksize[idx])
-        start = first(range) + p.blocksize[idx] * (subinds[idx] - 1)
-        stop = subinds[idx] == bsz + 1 ? last(range) : first(range) + p.blocksize[idx] * subinds[idx] - 1
-        start:stop
-    end
-    overlap = intersect(sd[1], sd[2])
-    if !isempty(overlap)
-        row_offset = first(sd[1]) - 1
-        col_offset = first(sd[2]) - 1
-        for g in overlap
-            sA[g - row_offset, g - col_offset] = 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, d, p), 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(p::BlocksOrAuto, s::UniformScaling, dims::Dims; assignment::AssignmentType = :arbitrary)
-    length(dims) == 2 || throw(ArgumentError("UniformScaling allocation requires exactly two dimensions"))
-    return DMatrix(p, s, dims; assignment)
-end
-DArray(p::BlocksOrAuto, s::UniformScaling, dims::Integer...; assignment::AssignmentType = :arbitrary) =
-    DArray(p, s, dims; assignment)
-DArray(::AutoBlocks, s::UniformScaling, dims::Dims; assignment::AssignmentType = :arbitrary) =
-    DArray(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

@@ -123,6 +123,7 @@ 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

@@ -67,9 +67,6 @@ 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])
@@ -390,8 +387,6 @@ function lookup_parts(A::DArray, ps::AbstractArray, subdmns::DomainBlocks{N}, d:
     return lookup_parts(A, ps, subdmns, ArrayDomain(new_d))
 end
 
-lookup_parts(A::DArray, ps::AbstractArray, subdmns::DomainBlocks{N}, d::CartesianIndices) where {N} = lookup_parts(A, ps, subdmns, ArrayDomain(d.indices))
-
 """
     Base.fetch(c::DArray)
 

src/array/indexing.jl

@@ -1,55 +1,5 @@
 ### getindex
 
-struct GetIndex{T,N} <: ArrayOp{T,N}
-    input::ArrayOp
-    idx::Tuple
-end
-
-GetIndex(input::ArrayOp, idx::Tuple) =
-    GetIndex{eltype(input), ndims(input)}(input, idx)
-
-function flatten(subdomains, subchunks, partitioning)
-    valdim = findfirst(j -> j != 1:1, subdomains[1].indexes)
-    flatc = []
-    flats = Array{ArrayDomain{1, Tuple{UnitRange{Int64}}}}(undef, 0)
-    map(x -> push!(flats, ArrayDomain(x.indexes[valdim])), subdomains)
-    map(x -> push!(flatc, x), subchunks)
-    newb = Blocks(partitioning.blocksize[valdim])
-    return flats, flatc, newb
-end
-
-function stage(ctx::Context, gidx::GetIndex)
-    inp = stage(ctx, gidx.input)
-
-    dmn = domain(inp)
-    idxs = [if isa(gidx.idx[i], Colon)
-        indexes(dmn)[i]
-    else
-        gidx.idx[i]
-    end for i in 1:length(gidx.idx)]
-
-    # Figure out output dimension
-    d = ArrayDomain(idxs)
-    subchunks, subdomains = Dagger.lookup_parts(inp, chunks(inp), domainchunks(inp), d)
-    d1 = alignfirst(d)
-    newb = inp.partitioning
-    if ndims(d1) != ndims(subdomains)
-        subdomains, subchunks, newb = flatten(subdomains, subchunks, inp.partitioning)
-    end
-    DArray(eltype(inp), d1, subdomains, subchunks, newb)
-end
-
-function size(x::GetIndex)
-    map(a -> a[2] isa Colon ?
-        size(x.input, a[1]) : length(a[2]),
-        enumerate(x.idx)) |> Tuple
-end
-
-Base.getindex(c::ArrayOp, idx::ArrayDomain) =
-    _to_darray(GetIndex(c, indexes(idx)))
-Base.getindex(c::ArrayOp, idx...) =
-    _to_darray(GetIndex(c, idx))
-
 const GETINDEX_CACHE = TaskLocalValue{Dict{Tuple,Any}}(()->Dict{Tuple,Any}())
 const GETINDEX_CACHE_SIZE = ScopedValue{Int}(0)
 with_index_caching(f, size::Integer=1) = with(f, GETINDEX_CACHE_SIZE=>size)
@@ -86,7 +36,6 @@ with_index_caching(f, size::Integer=1) = with(f, GETINDEX_CACHE_SIZE=>size)
     # Return the value
     return GPUArraysCore.@allowscalar 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)

src/array/linalg.jl

@@ -105,8 +105,10 @@ function LinearAlgebra.LAPACK.chkfinite(A::DArray)
 end
 
 DMatrix{T}(::LinearAlgebra.UniformScaling, m::Int, n::Int, IBlocks::Blocks) where T = DMatrix(Matrix{T}(I, m, n), IBlocks)
+DMatrix(::LinearAlgebra.UniformScaling{T}, m::Int, n::Int, IBlocks::Blocks) where T = DMatrix(Matrix{T}(I, m, n), IBlocks)
 
 DMatrix{T}(::LinearAlgebra.UniformScaling, size::Tuple, IBlocks::Blocks) where T = DMatrix(Matrix{T}(I, size), IBlocks)
+DMatrix(::LinearAlgebra.UniformScaling{T}, size::Tuple, IBlocks::Blocks) where T = DMatrix(Matrix{T}(I, size), IBlocks)
 
 function LinearAlgebra.inv(F::LU{T,<:DMatrix}) where T 
     n = size(F, 1)