Skip to content

Commit 55d83e7

Browse files
committed
Use :contiguous in N-index pyarray_offset
The Vararg method ignored the flag, and it is also the one picked for 1-D linear indexing, so sum over a contiguous vector didn't vectorise.
1 parent 0ec0320 commit 55d83e7

2 files changed

Lines changed: 19 additions & 1 deletion

File tree

‎src/Wrap/PyArray.jl‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -639,7 +639,13 @@ function pyarray_offset(x::PyArray{T,N,F}, i::Int) where {T,N,F}
639639
end
640640

641641
function pyarray_offset(x::PyArray{T,N,F}, i::Vararg{Int,N}) where {T,N,F}
642-
sum((i .- 1) .* x.strides)
642+
# Strides from the size make the unit first stride a compile-time constant, so loops vectorise.
643+
# Also matters for N == 1, where this method is picked over `pyarray_offset(x, i::Int)`.
644+
if (:contiguous in F)
645+
sum((i .- 1) .* Utils.size_to_fstrides(sizeof(pyarray_get_R(T)), x.size))
646+
else
647+
sum((i .- 1) .* x.strides)
648+
end
643649
end
644650

645651
function pyarray_load(::Type{T}, p::Ptr{R}) where {T,R}

‎test/Wrap.jl‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -100,6 +100,18 @@
100100
)
101101
@test_throws Exception PyArray(nd; array = false, buffer = true)
102102
end
103+
@testset "contiguous 2D" begin
104+
tb = pyimport("_testbuffer")
105+
nd = tb.ndarray(
106+
pylist(1:6),
107+
shape = pylist([2, 3]),
108+
format = "i",
109+
flags = tb.ND_FORTRAN,
110+
)
111+
a = PyArray(nd; array = false)
112+
@test a isa PyArray{Cint,2,(:linear,:contiguous)}
113+
@test [a[i, j] for i = 1:2, j = 1:3] == reshape(1:6, 2, 3)
114+
end
103115
end
104116

105117
@testitem "PyDict" begin

0 commit comments

Comments
 (0)