Skip to content

Commit 9698fa8

Browse files
authored
Add fallback gpu setindex (#137)
* Add fallback gpu setindex * Update ConcreteRArray.jl * Update ConcreteRArray.jl * Update ConcreteRArray.jl * Update ConcreteRArray.jl * Update ConcreteRArray.jl * Update ConcreteRArray.jl * Update ConcreteRArray.jl
1 parent f3c65db commit 9698fa8

File tree

1 file changed

+34
-1
lines changed

1 file changed

+34
-1
lines changed

src/ConcreteRArray.jl

Lines changed: 34 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -122,6 +122,7 @@ function Base.show(io::IO, X::ConcreteRArray)
122122
return Base.show(io, convert(Array, X))
123123
end
124124

125+
const getindex_warned = Ref(false)
125126
function Base.getindex(a::ConcreteRArray{T}, args::Vararg{Int,N}) where {T,N}
126127
if a.data == XLA.AsyncEmptyBuffer
127128
throw("Cannot getindex from empty buffer")
@@ -143,9 +144,26 @@ function Base.getindex(a::ConcreteRArray{T}, args::Vararg{Int,N}) where {T,N}
143144
return unsafe_load(ptr, start)
144145
end
145146
end
147+
if !getindex_warned[]
148+
@warn(
149+
"""Performing scalar get-indexing on task $(current_task()).
150+
Invocation resulted in scalar indexing of a ConcreteRArray.
151+
This is typically caused by calling an iterating implementation of a method.
152+
Such implementations *do not* execute on device, but very slowly on the CPU,
153+
and require expensive copies and synchronization each time and therefore should be avoided."""
154+
)
155+
getindex_warned[] = true
156+
end
146157
return convert(Array, a)[args...]
147158
end
148159

160+
function mysetindex!(a, v, args::Vararg{Int,N}) where {N}
161+
Base.setindex!(a, v, args...)
162+
nothing
163+
end
164+
165+
const setindex_warned = Ref(false)
166+
149167
function Base.setindex!(a::ConcreteRArray{T}, v, args::Vararg{Int,N}) where {T,N}
150168
if a.data == XLA.AsyncEmptyBuffer
151169
throw("Cannot setindex! to empty buffer")
@@ -167,7 +185,22 @@ function Base.setindex!(a::ConcreteRArray{T}, v, args::Vararg{Int,N}) where {T,N
167185
end
168186
return a
169187
end
170-
throw("Cannot setindex! to non-CPU buffer")
188+
if !setindex_warned[]
189+
@warn(
190+
"""Performing scalar set-indexing on task $(current_task()).
191+
Invocation resulted in scalar indexing of a ConcreteRArray.
192+
This is typically caused by calling an iterating implementation of a method.
193+
Such implementations *do not* execute on device, but very slowly on the CPU,
194+
and require expensive copies and synchronization each time and therefore should be avoided.
195+
196+
This error message will only be printed for the first invocation for brevity.
197+
"""
198+
)
199+
setindex_warned[] = true
200+
end
201+
fn = Reactant.compile(mysetindex!, (a, v, args...))
202+
fn(a, v, args...)
203+
return a
171204
end
172205

173206
# TODO is there any way to allocate an uninitialized buffer in XLA?

0 commit comments

Comments
 (0)