Skip to content

Commit 3b26ab2

Browse files
committed
Vectorise indexing of contiguous PyArrays
Branch on strides[1] == sizeof(R) so LLVM can version loops on the contiguous case. No change to the type parameters.
1 parent de68e38 commit 3b26ab2

3 files changed

Lines changed: 27 additions & 3 deletions

File tree

‎CHANGELOG.md‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
# Changelog
22

33
## Unreleased
4+
* Faster indexing into contiguous `PyArray`s.
45
* Bug fixes.
56

67
## 0.9.36 (2026-09-18)

‎src/Wrap/PyArray.jl‎

Lines changed: 9 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -672,11 +672,17 @@ end
672672
end
673673

674674
pyarray_offset(x::PyArray{T,N,M,true}, i::Int) where {T,N,M} =
675-
N == 0 ? 0 : (i - 1) * x.strides[1]
676-
pyarray_offset(x::PyArray{T,1,M,true}, i::Int) where {T,M} = (i - 1) .* x.strides[1]
677-
pyarray_offset(x::PyArray{T,N}, i::Vararg{Int,N}) where {T,N} = sum((i .- 1) .* x.strides)
675+
N == 0 ? 0 : pyarray_offset1(x, i)
676+
pyarray_offset(x::PyArray{T,1,M,true}, i::Int) where {T,M} = pyarray_offset1(x, i)
677+
pyarray_offset(x::PyArray{T,N}, i::Vararg{Int,N}) where {T,N} =
678+
pyarray_offset1(x, i[1]) + sum((Base.tail(i) .- 1) .* Base.tail(x.strides); init = 0)
678679
pyarray_offset(x::PyArray{T,0}) where {T} = 0
679680

681+
# The type can't say the first stride is the element size (strided arrays can be linear too).
682+
# Branching on it lets LLVM version loops on the contiguous case and vectorise them.
683+
pyarray_offset1(x::PyArray{T,N,M,L,R}, i::Int) where {T,N,M,L,R} =
684+
x.strides[1] == sizeof(R) ? (i - 1) * sizeof(R) : (i - 1) * x.strides[1]
685+
680686
function pyarray_load(::Type{T}, p::Ptr{R}) where {T,R}
681687
if R == T
682688
unsafe_load(p)

‎test/Wrap.jl‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -99,6 +99,23 @@
9999
)
100100
@test_throws Exception PyArray(nd; array = false, buffer = true)
101101
end
102+
@testset "linear with unit and non-unit first stride" begin
103+
tb = pyimport("_testbuffer")
104+
nd = tb.ndarray(
105+
pylist(1:24),
106+
shape = pylist([4, 6]),
107+
format = "i",
108+
flags = tb.ND_FORTRAN,
109+
)
110+
e = reshape(1:24, 4, 6)
111+
for (v, ev) in
112+
((nd, e), (nd[pyslice(nothing, nothing, 2), pyslice(nothing)], e[1:2:end, :]))
113+
a = PyArray(v; array = false)
114+
@test a isa PyArray{Cint,2,false,true,Cint}
115+
@test [a[i, j] for i in axes(a, 1), j in axes(a, 2)] == ev
116+
@test [a[i] for i in eachindex(IndexLinear(), a)] == vec(ev)
117+
end
118+
end
102119
end
103120

104121
@testitem "PyDict" begin

0 commit comments

Comments
 (0)