DArray: Indexing improvements and fixes

Adds optional task-local cache for getindex
Adds setindex! operation
Fixes lookup_parts indexing fallback for linear indexing
Adds allowscalar logic to prevent slow paths
Adds first/last helpers

Author: jpsamaroo

Diff for: Project.toml

@@ -20,6 +20,7 @@ SharedArrays = "1a1011a3-84de-559e-8e89-a11a2f7dc383"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
+TaskLocalValues = "ed4db957-447d-4319-bfb6-7fa9ae7ecf34"
 TimespanLogging = "a526e669-04d3-4846-9525-c66122c55f63"
 UUIDs = "cf7118a7-6976-5b1a-9a39-7adc72f591a4"
 
@@ -34,6 +35,7 @@ Requires = "1"
 ScopedValues = "1.1"
 Statistics = "1"
 StatsBase = "0.28, 0.29, 0.30, 0.31, 0.32, 0.33, 0.34"
+TaskLocalValues = "0.1"
 TimespanLogging = "0.1"
 julia = "1.8"
 

Diff for: docs/src/darray.md

@@ -134,12 +134,7 @@ julia> DZ = DY .* 3
 Dagger.DArray{Float64, 2, Blocks{2}, typeof(cat)}(100, 100)
 ```
 
-Now, `DZ` will contain the result of computing `(DX .+ DX) .* 3`. Note that
-`DArray` objects are immutable, and operations on them are thus functional
-transformations of their input `DArray`.
-
-!!! note
-    Support for mutation of `DArray`s is planned for a future release
+Now, `DZ` will contain the result of computing `(DX .+ DX) .* 3`.
 
 ```
 julia> Dagger.chunks(DZ)
@@ -208,14 +203,15 @@ julia> collect(DZ)
 ```
 
 A variety of other operations exist on the `DArray`, and it should generally
-behavior otherwise similar to any other `AbstractArray` type. If you find that
+behave otherwise similar to any other `AbstractArray` type. If you find that
 it's missing an operation that you need, please file an issue!
 
 ### Known Supported Operations
 
 This list is not exhaustive, but documents operations which are known to work well with the `DArray`:
 
 From `Base`:
+- `getindex`/`setindex!`
 - Broadcasting
 - `similar`/`copy`/`copyto!`
 - `map`/`reduce`/`mapreduce`

Diff for: src/Dagger.jl

@@ -66,7 +66,7 @@ include("array/copy.jl")
 include("file-io.jl")
 
 include("array/operators.jl")
-include("array/getindex.jl")
+include("array/indexing.jl")
 include("array/setindex.jl")
 include("array/matrix.jl")
 include("array/sparse_partition.jl")

Diff for: src/array/darray.jl

@@ -234,7 +234,7 @@ Base.:(/)(x::DArray{T,N,B,F}, y::U) where {T<:Real,U<:Real,N,B,F} =
 A `view` of a `DArray` chunk returns a `DArray` of `Thunk`s.
 """
 function Base.view(c::DArray, d)
-    subchunks, subdomains = lookup_parts(chunks(c), domainchunks(c), d)
+    subchunks, subdomains = lookup_parts(c, chunks(c), domainchunks(c), d)
     d1 = alignfirst(d)
     DArray(eltype(c), d1, subdomains, subchunks, c.partitioning, c.concat)
 end
@@ -272,7 +272,7 @@ function group_indices(cumlength, idxs::AbstractRange)
 end
 
 _cumsum(x::AbstractArray) = length(x) == 0 ? Int[] : cumsum(x)
-function lookup_parts(ps::AbstractArray, subdmns::DomainBlocks{N}, d::ArrayDomain{N}) where N
+function lookup_parts(A::DArray, ps::AbstractArray, subdmns::DomainBlocks{N}, d::ArrayDomain{N}) where N
     groups = map(group_indices, subdmns.cumlength, indexes(d))
     sz = map(length, groups)
     pieces = Array{Any}(undef, sz)
@@ -284,7 +284,15 @@ function lookup_parts(ps::AbstractArray, subdmns::DomainBlocks{N}, d::ArrayDomai
     end
     out_cumlength = map(g->_cumsum(map(x->length(x[2]), g)), groups)
     out_dmn = DomainBlocks(ntuple(x->1,Val(N)), out_cumlength)
-    pieces, out_dmn
+    return pieces, out_dmn
+end
+function lookup_parts(A::DArray, ps::AbstractArray, subdmns::DomainBlocks{N}, d::ArrayDomain{S}) 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))
 end
 
 """

Diff for: src/array/getindex.jl

@@ -0,0 +1,153 @@
+using TaskLocalValues
+
+### 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 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
+    view(inp, ArrayDomain(idxs))
+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)
+function Base.getindex(A::DArray{T,N}, idx::NTuple{N,Int}) where {T,N}
+    # Scalar indexing check
+    assert_allowscalar()
+
+    # Boundscheck
+    checkbounds(A, idx...)
+
+    # Find the associated partition and offset within it
+    part_idx, offset_idx = partition_for(A, idx)
+
+    # If the partition is cached, use that for lookup
+    cache = GETINDEX_CACHE[]
+    cache_size = GETINDEX_CACHE_SIZE[]
+    if cache_size > 0 && haskey(cache, part_idx)
+        return cache[part_idx][offset_idx...]
+    end
+
+    # Uncached, fetch the partition
+    part = fetch(A.chunks[part_idx...])
+
+    # Insert the partition into the cache
+    if cache_size > 0
+        if length(cache) >= cache_size
+            # Evict a random entry
+            key = rand(keys(cache))
+            delete!(cache, key)
+        end
+        cache[part_idx] = part
+    end
+
+    # 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)
+    for dim in 1:N
+        part_idx_slice = @view part_idx[1:(dim-1)]
+        trailing_idx_slice = ntuple(i->Colon(), N-dim)
+        sds = @view A.subdomains[part_idx_slice..., :, trailing_idx_slice...]
+        for (sd_idx, sd) in enumerate(sds)
+            sd_range = (sd.indexes::NTuple{N,UnitRange{Int}})[dim]
+            if sd_range.start <= idx[dim] <= sd_range.stop
+                part_idx[dim] = sd_idx
+                offset_idx[dim] = idx[dim] - sd_range.start + 1
+                break
+            end
+        end
+    end
+    return (part_idx...,), (offset_idx...,)
+end
+Base.getindex(A::DArray, idx::Integer...) =
+    getindex(A, idx)
+Base.getindex(A::DArray, idx::Integer) =
+    getindex(A, Base._ind2sub(A, idx))
+Base.getindex(A::DArray, idx::CartesianIndex) =
+    getindex(A, Tuple(idx))
+
+### setindex!
+
+function Base.setindex!(A::DArray{T,N}, value, idx::NTuple{N,Int}) where {T,N}
+    # Scalar indexing check
+    assert_allowscalar()
+
+    # Boundscheck
+    checkbounds(A, idx...)
+
+    # Find the associated partition and offset within it
+    part_idx, offset_idx = partition_for(A, idx)
+
+    # If the partition is cached, evict it
+    cache = GETINDEX_CACHE[]
+    if haskey(cache, part_idx)
+        delete!(cache, part_idx)
+    end
+
+    # Set the value
+    part = A.chunks[part_idx...]
+    space = memory_space(part)
+    scope = Dagger.scope(worker=root_worker_id(space))
+    return fetch(Dagger.@spawn scope=scope setindex!(part, value, offset_idx...))
+end
+Base.setindex!(A::DArray, value, idx::Integer...) =
+    setindex!(A, value, idx)
+Base.setindex!(A::DArray, value, idx::Integer) =
+    setindex!(A, value, Base._ind2sub(A, idx))
+Base.setindex!(A::DArray, value, idx::CartesianIndex) =
+    setindex!(A, value, Tuple(idx))
+
+### Allow/disallow scalar indexing
+
+const ALLOWSCALAR_TASK = TaskLocalValue{Bool}(()->true)
+const ALLOWSCALAR_SCOPE = ScopedValue{Bool}(false)
+isallowscalar() = ALLOWSCALAR_TASK[] || ALLOWSCALAR_SCOPE[]
+function assert_allowscalar()
+    if !isallowscalar()
+        throw(ArgumentError("Scalar indexing is disallowed\nSee `allowscalar` and `allowscalar!` for ways to disable this check, if necessary"))
+    end
+end
+"Allow/disallow scalar indexing for the current task."
+function allowscalar!(allow::Bool=true)
+    ALLOWSCALAR_TASK[] = allow
+end
+"Allow/disallow scalar indexing for the duration of executing `f`."
+function allowscalar(f, allow::Bool=true)
+    old = ALLOWSCALAR_TASK[]
+    allowscalar!(allow)
+    try
+        return with(f, ALLOWSCALAR_SCOPE=>allow)
+    finally
+        allowscalar!(old)
+    end
+end

Diff for: src/array/indexing.jl

@@ -113,3 +113,8 @@ function stage(ctx::Context, node::MapChunk)
     # TODO: Concrete type
     DArray(Any, domain(inputs[1]), domainchunks(inputs[1]), thunks)
 end
+
+# Basic indexing helpers
+
+Base.first(A::DArray) = A[begin]
+Base.last(A::DArray) = A[end]

Diff for: src/array/operators.jl

@@ -17,7 +17,7 @@ function getindex(x::DomainBlocks{N}, idx::Int) where N
     if N == 1
         _getindex(x, (idx,))
     else
-        _getindex(x, ind2sub(x, idx))
+        _getindex(x, Base._ind2sub(x, idx))
     end
 end
 

Diff for: src/lib/domain-blocks.jl

@@ -3,6 +3,7 @@ abstract type MemorySpace end
 struct CPURAMMemorySpace <: MemorySpace
     owner::Int
 end
+root_worker_id(space::CPURAMMemorySpace) = space.owner
 
 memory_space(x) = CPURAMMemorySpace(myid())
 function memory_space(x::Chunk)