fix: Specialized ReshapedArray dispatch to resolve setindex! ambiguities

Author: shreyas-omkar

src/host/indexing.jl

@@ -167,8 +167,14 @@ end
 function Base._unsafe_setindex!(::IndexStyle, A::WrappedGPUArray, x, Is::Vararg{Union{Real,AbstractArray}, N}) where N
     return vectorized_setindex!(A, x, Base.ensure_indexable(Is)...)
 end
-# And allow one more `ReshapedArray` wrapper to handle the `_maybe_reshape` optimization.
-function Base._unsafe_setindex!(::IndexStyle, A::Base.ReshapedArray{<:Any, <:Any, <:WrappedGPUArray}, x, Is::Vararg{Union{Real,AbstractArray}, N}) where N
+
+#Implementation for ReshapedArrays using Cartesian indexing to resolve dispatch ties.
+function Base._unsafe_setindex!(::Base.IndexCartesian, A::Base.ReshapedArray{T, N, <:WrappedGPUArray}, x, Is::Vararg{Union{Real, AbstractArray}, N}) where {T, N}
+    return vectorized_setindex!(A, x, Base.ensure_indexable(Is)...)
+end
+
+#Implementation for ReshapedArrays using Linear indexing to resolve dispatch ties.
+function Base._unsafe_setindex!(::Base.IndexLinear, A::Base.ReshapedArray{T, N, <:WrappedGPUArray}, x, Is::Vararg{Union{Real, AbstractArray}, N}) where {T, N}
     return vectorized_setindex!(A, x, Base.ensure_indexable(Is)...)
 end
 

test/testsuite/indexing.jl

@@ -284,3 +284,78 @@ end
         @test compare(argmin, AT, -rand(Int, 10))
     end
 end
+
+@testsuite "indexing combinatorial" (AT, eltypes) -> begin
+    @testset "Reshaped SubArray dispatch" for T in eltypes
+        #3D slice assignment
+        @testset "3D slice assignment" begin
+            let A = AT(ones(T, 4, 4, 4))
+                @views V = A[:, :, 1:2]
+                @allowscalar begin
+                    @test_nowarn V[:] .= T(0)
+                    @test all(Array(V) .== T(0))
+                end
+            end
+        end
+
+        #Logical indexing views
+        @testset "Logical mask view" begin
+            let A = AT(ones(T, 4, 4, 4))
+                mask = [true, false, true, false]
+                @views V = A[:, :, mask]
+                @allowscalar begin
+                    @test_nowarn V .+= T(2)
+                    @test all(Array(V) .== T(3))
+                end
+            end
+        end
+
+        #Nested Reshape layers
+        @testset "Nested Reshape" begin
+            let A = AT(ones(T, 4, 4, 4))
+                V = view(A, 1:2, 1:2, 1:2)
+                R1 = reshape(V, (4, 2))
+                R2 = reshape(R1, :)
+                @allowscalar begin
+                    @test_nowarn R2 .+= T(1)
+                    @test all(Array(R2) .== T(2))
+                end
+            end
+        end
+    end
+
+    @testset "Permuted and Reinterpreted Views" for T in eltypes
+        #PermutedDimsArray + Reshape
+        @testset "Reshaped PermutedDims" begin
+            let A = AT(ones(T, 4, 4))
+                P = PermutedDimsArray(A, (2, 1))
+                R = reshape(P, :)
+                @allowscalar begin
+                    @test_nowarn R[1:2] .= T(0)
+                    @test Array(R)[1] == T(0)
+                end
+            end
+        end
+
+        #Reinterpreted Reshape
+        @testset "Reshaped Reinterpreted" begin
+            let A = AT(ones(T, 4, 4))
+                IT = T <: Complex ? Complex{Int16} : Int16
+                R = reshape(reinterpret(IT, A), :)
+                @allowscalar begin
+                    @test_nowarn R[1] = IT(0)
+                    @test Array(R)[1] == IT(0)
+                end
+            end
+        end
+    end
+
+    @testset "Data parity with compare()" for T in eltypes
+        @test compare(AT, rand(T, 8, 8, 8)) do A
+            mask = [true, false, true, false, true, false, true, false]
+            @views V = A[:, mask, :]
+            @allowscalar V .+= T(1)
+            A
+        end
+    end
+end