Regression mapping over over GPU array and unit range

[sshin23]

On GPUArrays v10.3.1, map((x,y)->x+y, CUDA.zeros(2), 1:2) used to work without causing a scalar indexing issue. However, with the change in v11.1.0, it now throws the following error:

julia> map((x,y)->x+y, CUDA.zeros(2), 1:2)
ERROR: Scalar indexing is disallowed.
Invocation of getindex resulted in scalar indexing of a GPU array.
This is typically caused by calling an iterating implementation of a method.
Such implementations *do not* execute on the GPU, but very slowly on the CPU,
and therefore should be avoided.

If you want to allow scalar iteration, use `allowscalar` or `@allowscalar`
to enable scalar iteration globally or for the operations in question.
Stacktrace:
  [1] error(s::String)
    @ Base ./error.jl:35
  [2] errorscalar(op::String)
    @ GPUArraysCore ~/.julia/packages/GPUArraysCore/aNaXo/src/GPUArraysCore.jl:151
  [3] _assertscalar(op::String, behavior::GPUArraysCore.ScalarIndexing)
    @ GPUArraysCore ~/.julia/packages/GPUArraysCore/aNaXo/src/GPUArraysCore.jl:124
  [4] assertscalar(op::String)
    @ GPUArraysCore ~/.julia/packages/GPUArraysCore/aNaXo/src/GPUArraysCore.jl:112
  [5] getindex
    @ ~/.julia/packages/GPUArrays/sBzM5/src/host/indexing.jl:50 [inlined]
  [6] iterate
    @ ./abstractarray.jl:1209 [inlined]
  [7] iterate
    @ ./abstractarray.jl:1207 [inlined]
  [8] _zip_iterate_some
    @ ./iterators.jl:444 [inlined]
  [9] _zip_iterate_all
    @ ./iterators.jl:436 [inlined]
 [10] iterate
    @ ./iterators.jl:426 [inlined]
 [11] iterate
    @ ./generator.jl:45 [inlined]
 [12] collect(itr::Base.Generator{Base.Iterators.Zip{Tuple{CuArray{Float32, 1, CUDA.DeviceMemory}, UnitRange{Int64}}}, Base.var"#4#5"{var"#3#4"}})
    @ Base ./array.jl:791
 [13] map(f::Function, it::CuArray{Float32, 1, CUDA.DeviceMemory}, iters::UnitRange{Int64})
    @ Base ./abstractarray.jl:3495
 [14] top-level scope
    @ REPL[3]:1

I wonder if this was an intended change. CUDA.zeros(2) .+ (1:2) still seems to work on CUDA v5.6 without causing a scalar indexing issue.

cc: @frapac, @amontoison
MadNLP/MadNLP.jl#395

[maleadt]

This is caused by 8dfd805, which was done to resolve an ambiguity. Reinstating that functionality (and adding a test) is as simple as:

diff --git a/src/host/broadcast.jl b/src/host/broadcast.jl
index eaa50bb..6c703c4 100644
--- a/src/host/broadcast.jl
+++ b/src/host/broadcast.jl
@@ -82,8 +82,9 @@ end
 allequal(x) = true
 allequal(x, y, z...) = x == y && allequal(y, z...)

-function Base.map(f, xs::AnyGPUArray...)
+function Base.map(f, x::AnyGPUArray, xs::AbstractArray...)
     # if argument sizes match, their shape needs to be preserved
+    xs = (x, xs...)
     if allequal(size.(xs)...)
          return f.(xs...)
     end
@@ -103,7 +104,7 @@ function Base.map(f, xs::AnyGPUArray...)
     return map!(f, dest, xs...)
 end

-function Base.map!(f, dest::AnyGPUArray, xs::AnyGPUArray...)
+function Base.map!(f, dest::AnyGPUArray, xs::AbstractArray...)
     # custom broadcast, ignoring the container size mismatches
     # (avoids the reshape + view that our mapreduce impl has to do)
     indices = LinearIndices.((dest, xs...))
diff --git a/test/testsuite/broadcasting.jl b/test/testsuite/broadcasting.jl
index 81b028f..b6e852a 100644
--- a/test/testsuite/broadcasting.jl
+++ b/test/testsuite/broadcasting.jl
@@ -143,6 +143,9 @@ function broadcasting(AT, eltypes)
             @test compare(AT, rand(ET, 2,2), rand(ET, 2)) do x,y
                 map!(+, x, y)
             end
+            @test compare(AT, rand(ET, 2), 1:2) do x,y
+                map!(+, x, y)
+            end
         end

         @testset "map $ET" begin
@@ -155,6 +158,9 @@ function broadcasting(AT, eltypes)
             @test compare(AT, rand(ET, 2,2), rand(ET, 2)) do x,y
                 map(+, x, y)
             end
+            @test compare(AT, rand(ET, 2), 1:2) do x,y
+                map(+, x, y)
+            end
         end
     end

However, this introduces an ambiguity:

Error in testset JLArray/indexing multidimensional:
Error During Test at /home/tim/Julia/pkg/GPUArrays/test/testsuite/indexing.jl:108
  Test threw exception
  Expression: compare(getindex, AT, a, i')
  MethodError: map(::typeof(identity), ::LinearAlgebra.Adjoint{Bool, JLArray{Bool, 1}}) is ambiguous.

  Candidates:
    map(f, av::LinearAlgebra.Adjoint{T, <:AbstractVector} where T, avs::LinearAlgebra.Adjoint{T, <:AbstractVector} where T...)
      @ LinearAlgebra ~/Julia/depot/juliaup/julia-1.11.2+0.x64.linux.gnu/share/julia/stdlib/v1.11/LinearAlgebra/src/adjtrans.jl:399
    map(f, x::GPUArraysCore.AnyGPUArray, xs::AbstractArray...)
      @ GPUArrays ~/Julia/pkg/GPUArrays/src/host/broadcast.jl:85

I can't spend time on this right now, so feel free to take a stab.

[glwagner]

This also causes an issue over on CliMA/Oceananigans.jl#4036. However, I think our problem may be resolved by restoring previous behavior of map!, just glancing at LinearAlgebra it seems the ambiguity issue is only with map.