Slices and CuArrays

[mcabbott]

@MasonProtter points out that, somewhat to my surprise, mapslices-like things work on the GPU:

julia> using Zygote, TensorCast, CuArrays, BenchmarkTools
julia> CuArrays.allowscalar(false)
julia> let X = randn(10, 10, 10, 10, 10)
           Xcu = cu(X)
           f(A::AbstractArray{T, 4}) where {T} = reshape(sum(A, dims=(3,4)), size(A,1), size(A, 2))
           g(X) = @cast y[i, j, k] := f(X[:, :, :, :, k])[i, j]
           h_cu(a) = sum(g(a * Xcu))
           h_cpu(a)= sum(g(a * X))
           @show length(X)
           @btime $h_cu'(1)
           @btime $h_cpu'(1)
       end
length(X) = 100000
  1.660 ms (4676 allocations: 200.19 KiB)
  747.911 μs (374 allocations: 3.84 MiB)
351.3479618616162

julia> @pretty @cast y[i, j, k] := f(X[:, :, :, :, k])[i, j]
begin
    local armadillo = sliceview(X, (:, :, :, :, *))
    local porpoise = @__dot__(f(armadillo))
    local jaguar = red_glue(porpoise, (:, :, *))
    y = jaguar
end

It would be nice to (1) make sure this doesn't break, e.g. with #17, (2) understand which cases work or whether some don't, and (3) ideally make it fast?

[MasonProtter]

Maybe @maleadt has some thoughts on how we can set up CI to track this.

[mcabbott]

BTW, the sliced object here is an ordinary array of CuArrays. Which I think means that the iteration over slices happens on the CPU. I'm not sure precisely how that all works, but it sounds like it would work best for lots of work on few, large, slices.

julia> TensorCast.sliceview(cu(rand(2,3)), (:,*))
3-element Array{CuArray{Float32,1,CuArray{Float32,2,Nothing}},1}:
 Float32[0.61146635, 0.5048153]
 Float32[0.7750392, 0.87952733]
 Float32[0.7177145, 0.9034438] 

julia> TensorCast.red_glue(ans, (:,*)) # is just reduce(hcat, ans)
2×3 CuArray{Float32,2,Nothing}:
 0.611466  0.775039  0.717714
 0.504815  0.879527  0.903444

[maleadt]

GPU CI resources are documented here: https://github.com/JuliaGPU/gitlab-ci

[mcabbott]

Thanks, that doesn't sound too difficult.

[mcabbott]

This package still lacks its own GPU tests, sadly.

It could at least easily get fake GPU tests via JLArrays.

Besides slices above, it would be worth testing examples from #28 and #25.