Add support for Numpy 2.x by aMarcireau · Pull Request #429 · PyO3/rust-numpy

aMarcireau · 2024-06-04T00:26:03Z

The changes are based on recommendations from https://numpy.org/devdocs/numpy_2_0_migration_guide.html#c-api-changes.

The most visible user-facing change is the addition of two feature flags (numpy-1 and numpy-2). By default, both features are enabled and the code is compiled with support for both ABI versions (with runtime checks to select the right function offsets). Functions that are only available in numpy 1 or 2 are not exposed in this case. Disabling default features (for instance numpy = {version = "0.21.0", default-features = false, features = ["numpy-1"]}) exposes version-specific functions and fields but the library will panic if the runtime numpy version does not match.

I have not done much testing, this should be tried on different code bases (ideally ones that use low-level field access) before merging.

This currently uses std::sync::OnceLock to cache the runtime version. I realised too late that this is not compatible with the Minimum Supported Rust Version (it was introduced in 1.70.0). Using pyo3::sync::GILOnceCell isn't straightforward since py is not always available in functions that need to check the version to pick an implementation.

By default, the extension is compile with support for numpy 1 and 2 (with runtime checks to pick the right binary offset where needed). Features or fields that are specific to a version are hidden by default. Users can opt-out of numpy 1 + numpy 2 by disabling default features and selecting a version. The library panics if the runtime version does not match the compilation version if only one version is selected.

Essentially a port of https://github.com/numpy/numpy/blob/97356bc6f0d6538389a9eef475d883a0f4024c2a/numpy/_core/include/numpy/npy_2_compat.h

Cargo.toml

src/npyffi/mod.rs

src/npyffi/objects.rs

adamreichold · 2024-06-04T19:03:56Z

src/npyffi/mod.rs

+
+pub const NPY_2_0_API_VERSION: c_uint = 0x00000012;
+
+pub static ABI_API_VERSIONS: std::sync::OnceLock<(c_uint, c_uint)> = std::sync::OnceLock::new();


I think GILOnceCell is applicable here as we can just require that the accessor functions like PyDataType_FLAGS take a py: Python token. (We define them here so there is not need to conform exactly to the C signature just as there is no guarantee that they will stay in sync with the definitions in NumPy's C headers.)

Besides accessor functions, the version is also needed in the API functions that have a different offset in Numpy 1 and 2 (for instance PyArray_CopyInto) or only exist in Numpy 1 or 2 (using them with the wrong runtime version would otherwise result in memory corruption or a segfault).

I'm happy to switch to GILOnceCell but just wanted to check that changing higher-level function signatures was ok. For instance,

rust-numpy/src/dtype.rs

Line 259 in 2170e16

pub fn flags(&self) -> c_char {

will need py: Python so that it can call PyDataType_FLAGS).

If you have GIL ref &'py PyArray<..> or a bound ref Bound<'py, PyArray<..>>, this implies access to a GIL token via e.g. Bound::py. So I don't think this is an issue.

In any case, this will be a breaking release so we can change the API where necessary.

src/npyffi/objects.rs

adamreichold

Thank you for working on this! From looking at the code, I think we can take a simpler approach and always build with support for both NumPy 1.x and 2.x.

I also think GILOnceCell is appropriate to handle the version information.

adamreichold · 2024-06-04T19:10:54Z

src/npyffi/array.rs

+    #[cfg(all(feature = "numpy-1", not(feature = "numpy-2")))]
    impl_api![50; PyArray_CastTo(out: *mut PyArrayObject, mp: *mut PyArrayObject) -> c_int];
+    #[cfg(all(not(feature = "numpy-1"), feature = "numpy-2"))]
+    impl_api![50; PyArray_CopyInto(dst: *mut PyArrayObject, src: *mut PyArrayObject) -> c_int];


Since builds with support for versions will be common, I think this does not really add that much compile-time safety. I would rather suggest we extend the impl_api macro to do a runtime version check for function which are only available in one or the other version, e.g.

impl_api![@npy1 50; PyArray_CastTo(out: *mut PyArrayObject, mp: *mut PyArrayObject) -> c_int]; impl_api![@npy2 50; PyArray_CopyInto(dst: *mut PyArrayObject, src: *mut PyArrayObject) -> c_int];

That sounds good! I'll make these changes.

…propriate

aMarcireau · 2024-06-07T10:28:36Z

@adamreichold I just pushed a new version. Here's a brief overview of breaking changes.

The struct numpy::npyffi::PyArray_Descr now only contains public fields with the same position in numpy 1 and numpy 2. The remaining fields (elsize, alignment, metadata, subarray, names, fields, and c_metadata) may be accessed with the functions PyDataType_$field, where $field is ELSIZE, ALIGNMENT and so on.
The following functions of the struct numpy::PyArrayDescr and the trait numpy::PyArrayDescrMethods now have an extra parameter py: Python<'py> in first position: itemsize, alignment, flags, ndim, has_object, is_aligned_struct, has_subarray, and has_fields.
The type of dtype flags is now u64 to support numpy 2. It was c_char before . Negative values were technically possible but the numpy 1 to numpy 2 back port suggests to simply cast flag values to unsigned (C) chars, which is what we do here (in the accessor function).

Icxolu · 2024-06-07T18:30:47Z

The following functions of the struct numpy::PyArrayDescr and the trait numpy::PyArrayDescrMethods now have an extra parameter py: Python<'py> in first position: itemsize, alignment, flags, ndim, has_object, is_aligned_struct, has_subarray, and has_fields.

I don't think it is neccessary to change the API here. Since these are Python types we already have the proof that the GIL is held. You can just get the token via self.py(). For the trait you probably have to remove the default implementation and move them into the impl on Bound to get access to the token.

aMarcireau · 2024-06-08T01:13:49Z

@Icxolu Good point! I made the changes that your suggested.

@adamreichold Let me know if you would like me to squash the commits that modified src/dtype, since I ended up rolling back many of the edits.

adamreichold · 2024-06-18T18:46:13Z

Sorry for not getting to this yet. We are in crunch mode at $DAYJOB until next week. Will look into as soon as I can.

stinodego

@aMarcireau We tried building Polars using this branch, but there was a minor issue on Windows.

I forked your repo and with this one-line fix everything seems to work as expected:
https://github.com/stinodego/rust-numpy/commit/9ba9962ae57ba26e35babdce6f179edf5fe5b9c8

@stinodego

Contributed by @stinodego

aMarcireau · 2024-07-08T17:21:21Z

Thanks @stinodego. I replicated your changes in this PR.

adamreichold · 2024-07-21T12:36:43Z

src/dtype.rs

-        unsafe { !(*self.as_dtype_ptr()).subarray.is_null() }
-    }
+    ///
+    /// Equivalent to PyDataType_HASSUBARRAY(self)


Including this in the documentation is nice, but please turn it into a link in the NumPy docs. (Same for PyData_HASFIELDS.)

These don't appear to be documented anywhere. We could link to the macro definitions in the source if desired: https://github.com/numpy/numpy/blob/b3ddf2fd33232b8939f48c7c68a61c10257cd0c5/numpy/_core/include/numpy/ndarrayobject.h#L254-L255

adamreichold · 2024-07-21T12:38:31Z

src/dtype.rs

    ///
    /// [dtype-flags]: https://numpy.org/doc/stable/reference/generated/numpy.dtype.flags.html
-    pub fn flags(&self) -> c_char {
+    pub fn flags(&self) -> u64 {


Please expand the documentation to explain the type differences between NumPy 1.x and 2.x.

Added a brief note in #442 so indicate the flags field was expanded. Not sure if you wanted more than that.

adamreichold · 2024-07-21T12:42:51Z

src/dtype.rs

+    }
+
    fn base(&self) -> Bound<'py, PyArrayDescr> {
        if !self.has_subarray() {


Now that we need those accessor functions, I think it would be nice to use only one access, i.e. something like

let subarray = unsafe { PyDataType_SUBARRAY(self.py(), self.as_dtype_ptr()).as_ref() }; if let Some(subarray) = subarray { unsafe { Bound::from_borrowed_ptr(self.py(), subarray.base.cast()).downcast_into_unchecked() } } else { return 0; }

and similar in the other places below.

Done in #442.

…n bounds for the MSRV build.

adamreichold · 2024-07-21T12:55:58Z

src/npyffi/objects.rs


+#[repr(C)]
+#[derive(Copy, Clone)]
+pub struct _PyArray_DescrNumPy2 {


I would prefer a more consistent naming here, e.g. _PyArray_Descr_NumPy1 and _PyArray_Descr_NumPy2. (PyDataType_ISLEGACY can stay as it is.)

Done in #442.

adamreichold · 2024-07-21T12:56:24Z

src/npyffi/objects.rs

+#[allow(non_snake_case)]
+#[inline(always)]
+pub unsafe fn PyDataType_ISLEGACY(dtype: *const PyArray_Descr) -> bool {
+    (*dtype).type_num < NPY_TYPES::NPY_VSTRING as i32 && (*dtype).type_num >= 0


Please use as _ or as c_int here to match the type definition.

Done in #442.

adamreichold · 2024-07-21T12:59:24Z

src/npyffi/objects.rs

+}
+
+macro_rules! define_descr_accessor {
+    ($name:ident, $property:ident, $type:ty, $legacy_only:literal, $zero:expr) => {


Please rename $zero to $default or $fallback or something like that.

Done in #442.

adamreichold · 2024-07-21T13:00:41Z

src/npyffi/objects.rs

    pub meta: PyArray_DatetimeMetaData,
 }
+
+// npy_packed_static_string and npy_string_allocator are opaque pointers


Please tag the comment as FIXME(aMarcireau) or as FIXME(adamreichold) if you might be unable to follow up on this as we do elsewhere in the code base.

Done in #442.

adamreichold · 2024-07-21T13:04:12Z

The dtype::tests::test_dtype_names unit test appears to fail in the CI running NumPy 2. More generally, we might want to have at least one test matrix entry that explicitly builds against NumPy 1 to keep ourselves honest in supporting both major versions.

(I am not what the macOS issue about installing Python is but I would be surprised if it were actually related to this PR. @davidhewitt Is this something that already affected PyO3 and which we could shamelessly copy here?)

adamreichold · 2024-07-21T13:10:27Z

src/npyffi/mod.rs

+    [$offset: expr; ..=1.26; $fname: ident ($($arg: ident: $t: ty),* $(,)?) $(-> $ret: ty)?] => {
+        impl_api![$offset; ..=0x00000011; $fname($($arg : $t), *) $(-> $ret)*];
+    };
+    [$offset: expr; 2.0..; $fname: ident ($($arg: ident: $t: ty),* $(,)?) $(-> $ret: ty)?] => {
+        impl_api![$offset; 0x00000012..; $fname($($arg : $t), *) $(-> $ret)*];
+    };


The pattern syntax is nice, but are any but the two NumPy 1 and NumPy 2 variants actually used?

If not, please simplify this to have exactly two patterns here, one for NumPy 1 implementing the maximum version of 0x11 and one for NumPy implementing the minimum version of 0x12.

(I think this would also imply that we can loose api_version_to_numpy_version_range as we really only need to say "requires NumPy 1/2" which is a nice follow-up simplification.)

Done in #442.

adamreichold

So I made another pass and I think the basic approach is sound, but we need to fix and extend the CI and resolve various smaller things before this can land.

adamreichold · 2024-07-21T13:29:38Z

(I am not what the macOS issue about installing Python is but I would be surprised if it were actually related to this PR. @davidhewitt Is this something that already affected PyO3 and which we could shamelessly copy here?)

Going to tackle this separately in #436...

adamreichold · 2024-07-21T13:51:15Z

@aMarcireau I merged the necessary fixes so that the CI passes on the main branch now. You should rebase and then make sure to relax the various numpy<2 requirements I added when you add one CI job to explicitly test the older NumPy version. (I would suggest doing this on a Linux runner.)

JRRudy1 · 2024-08-04T03:19:32Z

src/npyffi/objects.rs

+#[repr(C)]
+#[derive(Clone, Copy)]
+pub struct npy_static_string {
+    size: usize,


Can you make the size and buf fields pub please? They are public in the numpy2 C-API, and we'll need access to them if we ever want to add support for variable-width string arrays

Done in #442.

davidhewitt · 2024-09-23T20:37:44Z

Continued by #442. Thank you all for the efforts in this PR and its follow-up!

aMarcireau added 3 commits June 4, 2024 09:53

Remove outdated 1.0 functions, add new 2.0 functions

fb7d512

Essentially a port of https://github.com/numpy/numpy/blob/97356bc6f0d6538389a9eef475d883a0f4024c2a/numpy/_core/include/numpy/npy_2_compat.h

Use the new "universal" access functions

17955d3

aMarcireau mentioned this pull request Jun 4, 2024

Support for Numpy 2 #409

Closed

Fix runtime tests

7a5a567