1
1
# host-level indexing
2
2
3
+ using Base: @propagate_inbounds
4
+
3
5
4
6
# indexing operators
5
7
@@ -12,29 +14,29 @@ vectorized_indices(Is...) = Val{true}()
12
14
# by only implementing `getindex(A, ::Int)` etc. this is difficult due to
13
15
# ambiguities with the vectorized method that can take any index type.
14
16
15
- Base . @propagate_inbounds Base. getindex (A:: AbstractGPUArray , Is... ) =
17
+ @propagate_inbounds Base. getindex (A:: AbstractGPUArray , Is... ) =
16
18
_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... ) =
18
20
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... ) =
20
22
vectorized_getindex (A, to_indices (A, Is)... )
21
23
22
- Base . @propagate_inbounds Base. setindex! (A:: AbstractGPUArray , v, Is... ) =
24
+ @propagate_inbounds Base. setindex! (A:: AbstractGPUArray , v, Is... ) =
23
25
_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... ) =
25
27
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... ) =
27
29
vectorized_setindex! (A, v, to_indices (A, Is)... )
28
30
29
31
# # scalar indexing
30
32
31
- function scalar_getindex (A:: AbstractGPUArray{T} , Is... ) where T
33
+ @propagate_inbounds function scalar_getindex (A:: AbstractGPUArray{T} , Is... ) where T
32
34
@boundscheck checkbounds (A, Is... )
33
35
I = Base. _to_linear_index (A, Is... )
34
36
getindex (A, I)
35
37
end
36
38
37
- function scalar_setindex! (A:: AbstractGPUArray{T} , v, Is... ) where T
39
+ @propagate_inbounds function scalar_setindex! (A:: AbstractGPUArray{T} , v, Is... ) where T
38
40
@boundscheck checkbounds (A, Is... )
39
41
I = Base. _to_linear_index (A, Is... )
40
42
setindex! (A, v, I)
43
45
# we still dispatch to `Base.getindex(a, ::Int)` etc so that there's a single method to
44
46
# override when a back-end (e.g. with unified memory) wants to allow scalar indexing.
45
47
46
- function Base. getindex (A:: AbstractGPUArray{T} , I:: Int ) where T
48
+ @propagate_inbounds function Base. getindex (A:: AbstractGPUArray{T} , I:: Int ) where T
47
49
@boundscheck checkbounds (A, I)
48
50
assertscalar (" getindex" )
49
51
x = Array {T} (undef, 1 )
50
52
copyto! (x, 1 , A, I, 1 )
51
53
return x[1 ]
52
54
end
53
55
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
55
57
@boundscheck checkbounds (A, I)
56
58
assertscalar (" setindex!" )
57
59
x = T[v]
61
63
62
64
# # vectorized indexing
63
65
64
- function vectorized_getindex! (dest:: AbstractGPUArray , src:: AbstractArray , Is... )
66
+ @propagate_inbounds function vectorized_getindex! (dest:: AbstractGPUArray ,
67
+ src:: AbstractArray , Is... )
65
68
any (isempty, Is) && return dest # indexing with empty array
66
69
idims = map (length, Is)
67
70
@@ -73,7 +76,7 @@ function vectorized_getindex!(dest::AbstractGPUArray, src::AbstractArray, Is...)
73
76
return dest
74
77
end
75
78
76
- function vectorized_getindex (src:: AbstractGPUArray , Is... )
79
+ @propagate_inbounds function vectorized_getindex (src:: AbstractGPUArray , Is... )
77
80
shape = Base. index_shape (Is... )
78
81
dest = similar (src, shape)
79
82
return vectorized_getindex! (dest, src, Is... )
91
94
end
92
95
end
93
96
94
- function vectorized_setindex! (dest:: AbstractArray , src, Is... )
97
+ @propagate_inbounds function vectorized_setindex! (dest:: AbstractArray , src, Is... )
95
98
isempty (Is) && return dest
96
99
idims = length .(Is)
97
100
len = prod (idims)
0 commit comments