Noncontiguous slicing on GPU

Author: mtfishman

src/abstractblocksparsearray/unblockedsubarray.jl

@@ -27,6 +27,27 @@ function TypeParameterAccessors.similartype(arraytype::Type{<:UnblockedSubArray}
   return similartype(blocktype(parenttype(arraytype)), elt)
 end
 
+function Base.similar(
+  a::UnblockedSubArray, elt::Type, axes::Tuple{Base.OneTo,Vararg{Base.OneTo}}
+)
+  return similar(similartype(blocktype(parenttype(a)), elt), axes)
+end
+function Base.similar(a::UnblockedSubArray, elt::Type, size::Tuple{Int,Vararg{Int}})
+  return similar(a, elt, Base.OneTo.(size))
+end
+
+function ArrayLayouts.sub_materialize(a::UnblockedSubArray)
+  a_cpu = adapt(Array, a)
+  a_cpu′ = similar(a_cpu)
+  a_cpu′ .= a_cpu
+  if typeof(a) === typeof(a_cpu)
+    return a_cpu′
+  end
+  a′ = similar(a)
+  a′ .= a_cpu′
+  return a′
+end
+
 function Base.map!(
   f, a_dest::AbstractArray, a_src1::UnblockedSubArray, a_src_rest::UnblockedSubArray...
 )

test/test_basics.jl

@@ -56,7 +56,6 @@ arrayts = (Array, JLArray)
     a[Block(1, 1)] = dev(randn(elt, 2, 2))
     a[Block(2, 2)] = dev(randn(elt, 3, 3))
     @test_broken a[:, [2, 4]]
-    @test_broken a[[3, 5], [2, 4]]
 
     # TODO: Fix this and turn it into a proper test.
     a = dev(BlockSparseArray{elt}(undef, [2, 3], [2, 3]))
@@ -713,6 +712,13 @@ arrayts = (Array, JLArray)
     @test a[Block(2, 2)[1:2, 2:3]] == b
     @test blockstoredlength(a) == 1
 
+    # Non-contiguous slicing.
+    a = dev(BlockSparseArray{elt}(undef, [2, 3], [2, 3]))
+    a[Block(1, 1)] = dev(randn(elt, 2, 2))
+    a[Block(2, 2)] = dev(randn(elt, 3, 3))
+    I = ([3, 5], [2, 4])
+    @test a[I...] == Array(a)[I...]
+
     a = BlockSparseArray{elt}(undef, [2, 3], [2, 3])
     @views for b in [Block(1, 1), Block(2, 2)]
       a[b] = randn(elt, size(a[b]))