Skip to content

Commit cb7d20f

Browse files
authored
Add additional inbounds propagation annotations. (#548)
1 parent e256516 commit cb7d20f

File tree

1 file changed

+16
-13
lines changed

1 file changed

+16
-13
lines changed

src/host/indexing.jl

Lines changed: 16 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,7 @@
11
# host-level indexing
22

3+
using Base: @propagate_inbounds
4+
35

46
# indexing operators
57

@@ -12,29 +14,29 @@ vectorized_indices(Is...) = Val{true}()
1214
# by only implementing `getindex(A, ::Int)` etc. this is difficult due to
1315
# ambiguities with the vectorized method that can take any index type.
1416

15-
Base.@propagate_inbounds Base.getindex(A::AbstractGPUArray, Is...) =
17+
@propagate_inbounds Base.getindex(A::AbstractGPUArray, Is...) =
1618
_getindex(vectorized_indices(Is...), A, to_indices(A, Is)...)
17-
Base.@propagate_inbounds _getindex(::Val{false}, A::AbstractGPUArray, Is...) =
19+
@propagate_inbounds _getindex(::Val{false}, A::AbstractGPUArray, Is...) =
1820
scalar_getindex(A, to_indices(A, Is)...)
19-
Base.@propagate_inbounds _getindex(::Val{true}, A::AbstractGPUArray, Is...) =
21+
@propagate_inbounds _getindex(::Val{true}, A::AbstractGPUArray, Is...) =
2022
vectorized_getindex(A, to_indices(A, Is)...)
2123

22-
Base.@propagate_inbounds Base.setindex!(A::AbstractGPUArray, v, Is...) =
24+
@propagate_inbounds Base.setindex!(A::AbstractGPUArray, v, Is...) =
2325
_setindex!(vectorized_indices(Is...), A, v, to_indices(A, Is)...)
24-
Base.@propagate_inbounds _setindex!(::Val{false}, A::AbstractGPUArray, v, Is...) =
26+
@propagate_inbounds _setindex!(::Val{false}, A::AbstractGPUArray, v, Is...) =
2527
scalar_setindex!(A, v, to_indices(A, Is)...)
26-
Base.@propagate_inbounds _setindex!(::Val{true}, A::AbstractGPUArray, v, Is...) =
28+
@propagate_inbounds _setindex!(::Val{true}, A::AbstractGPUArray, v, Is...) =
2729
vectorized_setindex!(A, v, to_indices(A, Is)...)
2830

2931
## scalar indexing
3032

31-
function scalar_getindex(A::AbstractGPUArray{T}, Is...) where T
33+
@propagate_inbounds function scalar_getindex(A::AbstractGPUArray{T}, Is...) where T
3234
@boundscheck checkbounds(A, Is...)
3335
I = Base._to_linear_index(A, Is...)
3436
getindex(A, I)
3537
end
3638

37-
function scalar_setindex!(A::AbstractGPUArray{T}, v, Is...) where T
39+
@propagate_inbounds function scalar_setindex!(A::AbstractGPUArray{T}, v, Is...) where T
3840
@boundscheck checkbounds(A, Is...)
3941
I = Base._to_linear_index(A, Is...)
4042
setindex!(A, v, I)
@@ -43,15 +45,15 @@ end
4345
# we still dispatch to `Base.getindex(a, ::Int)` etc so that there's a single method to
4446
# override when a back-end (e.g. with unified memory) wants to allow scalar indexing.
4547

46-
function Base.getindex(A::AbstractGPUArray{T}, I::Int) where T
48+
@propagate_inbounds function Base.getindex(A::AbstractGPUArray{T}, I::Int) where T
4749
@boundscheck checkbounds(A, I)
4850
assertscalar("getindex")
4951
x = Array{T}(undef, 1)
5052
copyto!(x, 1, A, I, 1)
5153
return x[1]
5254
end
5355

54-
function Base.setindex!(A::AbstractGPUArray{T}, v, I::Int) where T
56+
@propagate_inbounds function Base.setindex!(A::AbstractGPUArray{T}, v, I::Int) where T
5557
@boundscheck checkbounds(A, I)
5658
assertscalar("setindex!")
5759
x = T[v]
@@ -61,7 +63,8 @@ end
6163

6264
## vectorized indexing
6365

64-
function vectorized_getindex!(dest::AbstractGPUArray, src::AbstractArray, Is...)
66+
@propagate_inbounds function vectorized_getindex!(dest::AbstractGPUArray,
67+
src::AbstractArray, Is...)
6568
any(isempty, Is) && return dest # indexing with empty array
6669
idims = map(length, Is)
6770

@@ -73,7 +76,7 @@ function vectorized_getindex!(dest::AbstractGPUArray, src::AbstractArray, Is...)
7376
return dest
7477
end
7578

76-
function vectorized_getindex(src::AbstractGPUArray, Is...)
79+
@propagate_inbounds function vectorized_getindex(src::AbstractGPUArray, Is...)
7780
shape = Base.index_shape(Is...)
7881
dest = similar(src, shape)
7982
return vectorized_getindex!(dest, src, Is...)
@@ -91,7 +94,7 @@ end
9194
end
9295
end
9396

94-
function vectorized_setindex!(dest::AbstractArray, src, Is...)
97+
@propagate_inbounds function vectorized_setindex!(dest::AbstractArray, src, Is...)
9598
isempty(Is) && return dest
9699
idims = length.(Is)
97100
len = prod(idims)

0 commit comments

Comments
 (0)