Work around type-instability in broadcast with 32-bit axes.

Author: maleadt

src/broadcast.jl

@@ -19,3 +19,14 @@ BroadcastStyle(::CUDA.CuArrayStyle{N, B1},
 # allocation of output arrays
 Base.similar(bc::Broadcasted{CuArrayStyle{N,B}}, ::Type{T}, dims) where {T,N,B} =
     similar(CuArray{T,length(dims),B}, dims)
+
+# Base.Broadcast can't handle Int32 axes
+# XXX: not using a quirk, as constprop/irinterpret is crucial here
+# XXX: 1.11 uses to_index i nstead of CartesianIndex
+Base.@propagate_inbounds Broadcast.newindex(arg::AnyCuDeviceArray, I::CartesianIndex) = CartesianIndex(_newindex(axes(arg), I.I))
+Base.@propagate_inbounds Broadcast.newindex(arg::AnyCuDeviceArray, I::Integer) = CartesianIndex(_newindex(axes(arg), (I,)))
+Base.@propagate_inbounds _newindex(ax::Tuple, I::Tuple) = # XXX: upstream this?
+  (ifelse(length(ax[1]) == 1, promote(ax[1][1], I[1])...), _newindex(Base.tail(ax), Base.tail(I))...)
+Base.@propagate_inbounds _newindex(ax::Tuple{}, I::Tuple) = ()
+Base.@propagate_inbounds _newindex(ax::Tuple, I::Tuple{}) = (ax[1][1], _newindex(Base.tail(ax), ())...)
+Base.@propagate_inbounds _newindex(ax::Tuple{}, I::Tuple{}) = ()

src/device/array.jl

@@ -43,6 +43,12 @@ end
 const CuDeviceVector = CuDeviceArray{T,1,A} where {T,A}
 const CuDeviceMatrix = CuDeviceArray{T,2,A} where {T,A}
 
+# anything that's (secretly) backed by a CuArray
+const AnyCuDeviceArray{T,N} = Union{CuDeviceArray{T,N}, WrappedArray{T,N,CuDeviceArray,CuDeviceArray{T,N}}}
+const AnyCuDeviceVector{T} = AnyCuDeviceArray{T,1}
+const AnyCuDeviceMatrix{T} = AnyCuDeviceArray{T,2}
+const AnyCuDeviceVecOrMat{T} = Union{AnyCuDeviceVector{T}, AnyCuDeviceMatrix{T}}
+
 
 ## array interface