From 55d83e7c3f1bcb228f2be2425e811cfdb4d85d23 Mon Sep 17 00:00:00 2001 From: Beforerr Date: Mon, 5 Oct 2026 18:15:34 -0600 Subject: [PATCH] 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. --- src/Wrap/PyArray.jl | 8 +++++++- test/Wrap.jl | 12 ++++++++++++ 2 files changed, 19 insertions(+), 1 deletion(-) diff --git a/src/Wrap/PyArray.jl b/src/Wrap/PyArray.jl index 64bb3176..95c639f1 100644 --- a/src/Wrap/PyArray.jl +++ b/src/Wrap/PyArray.jl @@ -639,7 +639,13 @@ function pyarray_offset(x::PyArray{T,N,F}, i::Int) where {T,N,F} end function pyarray_offset(x::PyArray{T,N,F}, i::Vararg{Int,N}) where {T,N,F} - sum((i .- 1) .* x.strides) + # Strides from the size make the unit first stride a compile-time constant, so loops vectorise. + # Also matters for N == 1, where this method is picked over `pyarray_offset(x, i::Int)`. + if (:contiguous in F) + sum((i .- 1) .* Utils.size_to_fstrides(sizeof(pyarray_get_R(T)), x.size)) + else + sum((i .- 1) .* x.strides) + end end function pyarray_load(::Type{T}, p::Ptr{R}) where {T,R} diff --git a/test/Wrap.jl b/test/Wrap.jl index 4720cdd4..38c7a40a 100644 --- a/test/Wrap.jl +++ b/test/Wrap.jl @@ -100,6 +100,18 @@ ) @test_throws Exception PyArray(nd; array = false, buffer = true) end + @testset "contiguous 2D" begin + tb = pyimport("_testbuffer") + nd = tb.ndarray( + pylist(1:6), + shape = pylist([2, 3]), + format = "i", + flags = tb.ND_FORTRAN, + ) + a = PyArray(nd; array = false) + @test a isa PyArray{Cint,2,(:linear,:contiguous)} + @test [a[i, j] for i = 1:2, j = 1:3] == reshape(1:6, 2, 3) + end end @testitem "PyDict" begin