Move to CartesianIndices alltogether.

Author: maleadt

src/device/indexing.jl

@@ -53,40 +53,6 @@ macro linearidx(A, ctxsym = :ctx)
     end
 end
 
-
-# Base functions that are sadly not fit for the the GPU yet (they only work for Int64)
-Base.@pure @inline function gpu_ind2sub(A::AbstractArray, ind::T) where T
-    _ind2sub(size(A), ind - T(1))
-end
-Base.@pure @inline function gpu_ind2sub(dims::NTuple{N}, ind::T) where {N, T}
-    _ind2sub(NTuple{N, T}(dims), ind - T(1))
-end
-Base.@pure @inline _ind2sub(::Tuple{}, ind::T) where {T} = (ind + T(1),)
-Base.@pure @inline function _ind2sub(indslast::NTuple{1}, ind::T) where T
-    ((ind + T(1)),)
-end
-Base.@pure @inline function _ind2sub(inds, ind::T) where T
-    r1 = inds[1]
-    indnext = div(ind, r1)
-    f = T(1); l = r1
-    (ind-l*indnext+f, _ind2sub(Base.tail(inds), indnext)...)
-end
-
-Base.@pure function gpu_sub2ind(dims::NTuple{N}, I::NTuple{N2, T}) where {N, N2, T}
-    Base.@_inline_meta
-    _sub2ind(NTuple{N, T}(dims), T(1), T(1), I...)
-end
-_sub2ind(x, L, ind) = ind
-function _sub2ind(::Tuple{}, L, ind, i::T, I::T...) where T
-    Base.@_inline_meta
-    ind + (i - T(1)) * L
-end
-function _sub2ind(inds, L, ind, i::IT, I::IT...) where IT
-    Base.@_inline_meta
-    r1 = inds[1]
-    _sub2ind(Base.tail(inds), L * r1, ind + (i - IT(1)) * L, I...)
-end
-
 """
     cartesianidx(A, ctxsym = :ctx)
 

src/host/abstractarray.jl

@@ -139,14 +139,12 @@ end
 Base.copyto!(dest::AbstractGPUArray, src::AbstractGPUArray) =
     copyto!(dest, CartesianIndices(dest), src, CartesianIndices(src))
 
-function copy_kernel!(ctx::AbstractKernelContext, dest, dest_offsets, src, src_offsets, shape, shape_dest, shape_source, length)
+function copy_kernel!(ctx::AbstractKernelContext, dest, dest_offsets, src, src_offsets, shape, length)
     i = linear_index(ctx)
     if i <= length
         # TODO can this be done faster and smarter?
-        idx = gpu_ind2sub(shape, i)
-        dest_idx = gpu_sub2ind(shape_dest, idx .+ dest_offsets)
-        src_idx = gpu_sub2ind(shape_source, idx .+ src_offsets)
-        @inbounds dest[dest_idx] = src[src_idx]
+        idx = CartesianIndices(shape)[i]
+        @inbounds dest[idx + dest_offsets] = src[idx + src_offsets]
     end
     return
 end
@@ -159,10 +157,10 @@ function Base.copyto!(dest::AbstractGPUArray{T, N}, destcrange::CartesianIndices
     end
     len = length(destcrange)
 
-    dest_offsets = first.(destcrange.indices) .- 1
-    src_offsets = first.(srccrange.indices) .- 1
+    dest_offsets = first(destcrange) - one(CartesianIndex{N})
+    src_offsets = first(srccrange) - one(CartesianIndex{N})
     gpu_call(copy_kernel!,
-             dest, dest_offsets, src, src_offsets, shape, size(dest), size(src), len;
+             dest, dest_offsets, src, src_offsets, shape, len;
              total_threads=len)
     dest
 end

src/host/indexing.jl

@@ -84,9 +84,8 @@ to_index(a, x::Base.LogicalIndex) = error("Logical indexing not implemented")
 @generated function index_kernel(ctx::AbstractKernelContext, dest::AbstractArray, src::AbstractArray, idims, Is)
     N = length(Is.parameters)
     quote
-        i = linear_index(ctx)
-        i > length(dest) && return
-        is = gpu_ind2sub(idims, i)
+        i = @linearidx dest
+        is = CartesianIndices(idims)[i]
         @nexprs $N i -> @inbounds I_i = Is[i][is[i]]
         @inbounds dest[i] = @ncall $N getindex src i -> I_i
         return
@@ -112,7 +111,7 @@ end
     quote
         i = linear_index(ctx)
         i > len && return
-        is = gpu_ind2sub(idims, i)
+        is = CartesianIndices(idims)[i]
         @inbounds setindex!(dest, bgetindex(src, i), $(idx...))
         return
     end

test/testsuite/base.jl

@@ -1,7 +1,6 @@
 function cartesian_iter(state, res, A, Asize)
     for i in CartesianIndices(Asize)
-        idx = GPUArrays.gpu_sub2ind(Asize, i.I)
-        res[idx] = A[idx]
+        res[i] = A[i]
     end
     return
 end