Optimize nth and nth_back for BoundListIterator #4810

Owen-CH-Leung · 2024-12-21T06:48:36Z

This PR optimizes nth and nth_back for BoundListIterator, and added unittest & benchmarks for these 2 APIs. Here are the benchmark of the optimized nth and nth_back:

On Stable toolchain

With Optimization

list_nth                time:   [351.21 ns 351.71 ns 352.28 ns]
                        change: [-88.294% -88.189% -88.089%] (p = 0.00 < 0.05)
                        Performance has improved.
Found 14 outliers among 100 measurements (14.00%)
  9 (9.00%) high mild
  5 (5.00%) high severe

list_nth_back           time:   [355.34 ns 356.09 ns 356.95 ns]
                        change: [-90.384% -90.343% -90.300%] (p = 0.00 < 0.05)
                        Performance has improved.
Found 7 outliers among 100 measurements (7.00%)
  5 (5.00%) high mild
  2 (2.00%) high severe

The default nth and nth_back implementation

list_nth                time:   [2.8422 µs 2.8502 µs 2.8592 µs]
                        change: [+706.38% +709.27% +712.94%] (p = 0.00 < 0.05)
                        Performance has regressed.
Found 11 outliers among 100 measurements (11.00%)
  6 (6.00%) high mild
  5 (5.00%) high severe

list_nth_back           time:   [3.6509 µs 3.6557 µs 3.6611 µs]
                        change: [+919.80% +923.58% +927.06%] (p = 0.00 < 0.05)
                        Performance has regressed.
Found 12 outliers among 100 measurements (12.00%)
  4 (4.00%) high mild
  8 (8.00%) high severe

On nightly toolchain

With Optimization

list_nth                time:   [290.88 ns 291.70 ns 292.71 ns]
                        change: [-0.5501% -0.2102% +0.2305%] (p = 0.27 > 0.05)
                        No change in performance detected.
Found 10 outliers among 100 measurements (10.00%)
  8 (8.00%) high mild
  2 (2.00%) high severe

list_nth_back           time:   [295.06 ns 296.29 ns 297.90 ns]
                        change: [-1.2387% -0.9560% -0.6230%] (p = 0.00 < 0.05)
                        Change within noise threshold.
Found 7 outliers among 100 measurements (7.00%)
  3 (3.00%) high mild
  4 (4.00%) high severe

The default nth and nth_back implementation

…nchmarks

davidhewitt · 2024-12-21T09:10:52Z

Awesome, thanks! Will seek to review soon. I guess we can do the same thing for tuple iteration?

Owen-CH-Leung · 2024-12-21T15:48:58Z

Awesome, thanks! Will seek to review soon. I guess we can do the same thing for tuple iteration?

Yup - I'll file a PR to optimise tuple iterator after the holiday =)

davidhewitt · 2024-12-24T11:04:55Z

src/types/list.rs

+        let length = self.length.min(self.list.len());
+        let target_index = self.index + n;
+        if self.index + n < length {
+            let item = unsafe { self.get_item(target_index) };


I wonder, is there a time-of-check to time-of-use bug here on the length? Not one for this PR, but a follow up I will try not to forget...

Yup I think that's a potential TOCTOU bug here. This particular implementation assumes the user ensures proper synchronization if they intend to use the iterator in a multi-threaded or mutable environment.

I think the current implementation of next also has a potential TOCTOU bug:

https://github.com/PyO3/pyo3/blob/main/src/types/list.rs#L495-L498

ngoldbaum · 2025-01-06T19:00:30Z

Thanks for working on this! Hopefully my PR should be merged in the next few days and then you can rebase this.

ngoldbaum · 2025-01-06T19:02:40Z

It might also make sense to implement advance_by on nightly. nth is implemented in terms of advance_by, so this improvement would fall out "for free" on nightly if you implement that function.

ngoldbaum

See inline comments, I think there are some issues in how this is set up for the limited API. I also think we should define nth and nth_back for all non-nightly builds. I also don't think it's necessary to call get_item in advance_by, per my reading of the advance_by docs.

ngoldbaum · 2025-01-09T16:04:00Z

src/types/list.rs

+            let length = length.0.min(list.len());
+            let target_index = index.0 + n;
+            if index.0 + n < length {
+                let item = list.get_item(target_index);


I don't think we need to actually do the get_item or check for an Err value, we only need to advance the iterator, which just requires the index.0 + n < length check.

ngoldbaum · 2025-01-09T16:06:43Z

src/types/list.rs

@@ -750,6 +807,27 @@ impl<'py> Iterator for BoundListIterator<'py> {
            None
        })
    }
+
+    #[inline]
+    #[cfg(all(Py_GIL_DISABLED, feature = "nightly"))]


Suggested change

#[cfg(all(Py_GIL_DISABLED, feature = "nightly"))]

#[cfg(feature = "nightly")]

I think the critical section is necessary though, since another thread might try to use the list while you're in advance_by. You don't need to be in a Py_GIL_DISABLED cfg block to use a critical section, it's a no-op on the GIL-enabled build.

ngoldbaum · 2025-01-09T16:07:18Z

src/types/list.rs

@@ -772,6 +850,14 @@ impl DoubleEndedIterator for BoundListIterator<'_> {
        }
    }

+    #[inline]
+    #[cfg(all(Py_GIL_DISABLED, not(feature = "nightly")))]


Suggested change

#[cfg(all(Py_GIL_DISABLED, not(feature = "nightly")))]

#[cfg(feature = "nightly")]

ngoldbaum · 2025-01-09T16:08:13Z

src/types/list.rs

@@ -625,6 +674,14 @@ impl<'py> Iterator for BoundListIterator<'py> {
        }
    }

+    #[inline]
+    #[cfg(all(Py_GIL_DISABLED, not(feature = "nightly")))]


Suggested change

#[cfg(all(Py_GIL_DISABLED, not(feature = "nightly")))]

#[cfg(feature = "nightly")]

ngoldbaum · 2025-01-09T16:10:49Z

src/types/list.rs

+    /// access to the list by holding a lock or by holding the innermost
+    /// critical section on the list.
+    #[inline]
+    #[cfg(not(Py_LIMITED_API))]


you should also add branches to nth and nth_back that use get_item instead of get_item_unchecked on the limited API

ngoldbaum · 2025-01-09T16:11:56Z

pyo3-benches/benches/bench_list.rs

@@ -66,6 +92,8 @@ fn sequence_from_list(b: &mut Bencher<'_>) {
 fn criterion_benchmark(c: &mut Criterion) {
    c.bench_function("iter_list", iter_list);
    c.bench_function("list_new", list_new);
+    c.bench_function("list_nth", list_nth);
+    c.bench_function("list_nth_back", list_nth_back);


Totally missed this when I added the new free-threaded implementations for folding operations in list. I'll add benchmarks for those in a separate PR.

ngoldbaum · 2025-01-09T16:12:42Z

noxfile.py

changes in this file seem unrelated, maybe send in a separate PR so we can get these ruff fixes merged quickly?

yeah - I noticed that ruff has released 0.9.0 2 hours ago and the CI didn't pin the ruff version, hence I hit this error which doesn't exist when the latest MR is merged.

I'll file another MR to quickly format this file

Owen-CH-Leung added 3 commits December 21, 2024 14:43

Optimize nth and nth_back for BoundListIterator. Add unit test and be…

cc9cabd

…nchmarks

Fix fmt and newsfragment CI

3a0c196

Fix clippy and changelog CI

40d38f3

Owen-CH-Leung marked this pull request as ready for review December 21, 2024 15:47

davidhewitt reviewed Dec 24, 2024

View reviewed changes

davidhewitt mentioned this pull request Jan 3, 2025

Implement locked iteration for PyList #4789

Merged

Owen-CH-Leung added 4 commits January 9, 2025 14:50

Merge branch 'main' into optimize_nth_and_nthback_for_boundlistiter

1b19616

Revise Impl of nth and nth_back. Impl advance_by

f6e95a8

Fix failing fmt

b0c749b

Fix failing ruff test

b2bf973

ngoldbaum reviewed Jan 9, 2025

View reviewed changes

Owen-CH-Leung mentioned this pull request Jan 9, 2025

ci: updates for Rust 1.84 #4846

Merged

Owen-CH-Leung added 3 commits January 10, 2025 08:52

Merge branch 'main' into optimize_nth_and_nthback_for_boundlistiter

4e86709

branch out nth, nth_unchecked, nth_back, nth_back_unchecked.

e4269c2

Fix fmt

6e18229

Owen-CH-Leung marked this pull request as draft January 10, 2025 08:22

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Optimize nth and nth_back for BoundListIterator #4810

Optimize nth and nth_back for BoundListIterator #4810

Owen-CH-Leung commented Dec 21, 2024 •

edited

Loading

davidhewitt commented Dec 21, 2024

Owen-CH-Leung commented Dec 21, 2024 •

edited

Loading

davidhewitt Dec 24, 2024

Owen-CH-Leung Dec 29, 2024 •

edited

Loading

Owen-CH-Leung Dec 29, 2024 •

edited

Loading

ngoldbaum commented Jan 6, 2025

ngoldbaum commented Jan 6, 2025

ngoldbaum left a comment

ngoldbaum Jan 9, 2025

ngoldbaum Jan 9, 2025

ngoldbaum Jan 9, 2025

ngoldbaum Jan 9, 2025

ngoldbaum Jan 9, 2025

ngoldbaum Jan 9, 2025

ngoldbaum Jan 9, 2025

Owen-CH-Leung Jan 9, 2025

	#[cfg(all(Py_GIL_DISABLED, feature = "nightly"))]
	#[cfg(feature = "nightly")]

	#[cfg(all(Py_GIL_DISABLED, not(feature = "nightly")))]
	#[cfg(feature = "nightly")]

Optimize nth and nth_back for BoundListIterator #4810

Are you sure you want to change the base?

Optimize nth and nth_back for BoundListIterator #4810

Conversation

Owen-CH-Leung commented Dec 21, 2024 • edited Loading

On Stable toolchain

On nightly toolchain

davidhewitt commented Dec 21, 2024

Owen-CH-Leung commented Dec 21, 2024 • edited Loading

Choose a reason for hiding this comment

Owen-CH-Leung Dec 29, 2024 • edited Loading

Choose a reason for hiding this comment

Owen-CH-Leung Dec 29, 2024 • edited Loading

Choose a reason for hiding this comment

ngoldbaum commented Jan 6, 2025

ngoldbaum commented Jan 6, 2025

ngoldbaum left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Owen-CH-Leung commented Dec 21, 2024 •

edited

Loading

Owen-CH-Leung commented Dec 21, 2024 •

edited

Loading

Owen-CH-Leung Dec 29, 2024 •

edited

Loading

Owen-CH-Leung Dec 29, 2024 •

edited

Loading