Merge pull request #251 from PyO3/parallel-example

Add new example showing how to use ndarray's parallel and blas-src features

Author: adamreichold

.github/workflows/ci.yml

@@ -20,17 +20,14 @@ jobs:
           default: true
       - uses: Swatinem/rust-cache@v1
         continue-on-error: true
-      - env:
-          CLIPPYFLAGS: --deny warnings --allow clippy::needless-lifetimes
-        run: |
+      - run: |
           cargo fmt --all -- --check
-          cargo clippy --tests -- $CLIPPYFLAGS
-          for example in examples/*; do (cd $example/; cargo clippy -- $CLIPPYFLAGS) || exit 1; done
+          cargo clippy --workspace --tests -- --deny warnings
 
   test:
     name: python${{ matrix.python-version }}-${{ matrix.platform.python-architecture }} ${{ matrix.platform.os }}
     runs-on: ${{ matrix.platform.os }}
-    needs: [lint, check-msrv, linalg-example]
+    needs: [lint, check-msrv, examples]
     strategy:
       fail-fast: false
       matrix:
@@ -84,8 +81,7 @@ jobs:
       - name: Test example
         run: |
           pip install tox
-          tox
-        working-directory: examples/simple-extension
+          tox -c examples/simple-extension
     env:
       CARGO_TERM_VERBOSE: true
       CARGO_BUILD_TARGET: ${{ matrix.platform.rust-target }}
@@ -143,7 +139,7 @@ jobs:
           tox
         working-directory: examples/simple-extension
 
-  linalg-example:
+  examples:
     runs-on: ubuntu-latest
     steps:
       - uses: actions/checkout@v2
@@ -161,8 +157,8 @@ jobs:
           default: true
       - uses: Swatinem/rust-cache@v1
         continue-on-error: true
-      - name: Test example
+      - name: Test examples
         run: |
           pip install tox
-          tox
-        working-directory: examples/linalg
+          tox -c examples/linalg
+          tox -c examples/parallel

examples/linalg/README.md

@@ -2,8 +2,7 @@
 
 An example extension with [ndarray-linalg](https://github.com/rust-ndarray/ndarray-linalg).
 
-Needs a Fortran compiler (e.g., `gfortran`) for building.
+Will link against a system-provided OpenBLAS.
 
 See [simple-extension's README](https://github.com/PyO3/rust-numpy/blob/main/examples/simple-extension/README.md)
 for an introduction.
-

examples/linalg/src/lib.rs

@@ -1,17 +1,16 @@
 use ndarray_linalg::solve::Inverse;
 use numpy::{IntoPyArray, PyArray2, PyReadonlyArray2};
-use pyo3::exceptions::PyRuntimeError;
-use pyo3::prelude::{pymodule, PyErr, PyModule, PyResult, Python};
+use pyo3::{exceptions::PyRuntimeError, pymodule, types::PyModule, PyErr, PyResult, Python};
 
 #[pymodule]
 fn rust_linalg(_py: Python<'_>, m: &PyModule) -> PyResult<()> {
     #[pyfn(m)]
     fn inv<'py>(py: Python<'py>, x: PyReadonlyArray2<'py, f64>) -> PyResult<&'py PyArray2<f64>> {
-        let x = x
-            .as_array()
+        let x = x.as_array();
+        let y = x
             .inv()
             .map_err(|e| PyErr::new::<PyRuntimeError, _>(format!("[rust_linalg] {}", e)))?;
-        Ok(x.into_pyarray(py))
+        Ok(y.into_pyarray(py))
     }
     Ok(())
 }

examples/linalg/tox.ini

@@ -7,5 +7,5 @@ deps =
     numpy
     pytest
 commands =
-    pip install .
+    pip install . -v
     pytest {posargs}

examples/parallel/Cargo.toml

@@ -0,0 +1,16 @@
+[package]
+name = "numpy-parallel-example"
+version = "0.1.0"
+authors = ["Yuji Kanagawa <[email protected]>"]
+edition = "2018"
+
+[lib]
+name = "rust_parallel"
+crate-type = ["cdylib"]
+
+[dependencies]
+pyo3 = { version = "0.15", features = ["extension-module"] }
+numpy = { path = "../.." }
+ndarray = { version = "0.15", features = ["rayon", "blas"] }
+blas-src = { version = "0.8", features = ["openblas"] }
+openblas-src = { version = "0.10", features = ["cblas", "system"] }

examples/parallel/README.md

@@ -0,0 +1,7 @@
+# rust-numpy example extension using optional ndarray features
+
+An example extension using [optional ndarray features](https://docs.rs/ndarray/latest/ndarray/doc/crate_feature_flags/index.html), parallel execution using Rayon and optimized kernels using BLAS in this case.
+
+See [simple-extension's README](https://github.com/PyO3/rust-numpy/blob/main/examples/simple-extension/README.md)
+for an introduction.
+

examples/parallel/pyproject.toml

@@ -0,0 +1,3 @@
+[build-system]
+build-backend = "maturin"
+requires = ["maturin>=0.12,<0.13"]

examples/parallel/src/lib.rs

@@ -0,0 +1,22 @@
+// We need to link `blas_src` directly, c.f. https://github.com/rust-ndarray/ndarray#how-to-enable-blas-integration
+extern crate blas_src;
+
+use ndarray::Zip;
+use numpy::{IntoPyArray, PyArray1, PyReadonlyArray1, PyReadonlyArray2};
+use pyo3::{pymodule, types::PyModule, PyResult, Python};
+
+#[pymodule]
+fn rust_parallel(_py: Python<'_>, m: &PyModule) -> PyResult<()> {
+    #[pyfn(m)]
+    fn rows_dot<'py>(
+        py: Python<'py>,
+        x: PyReadonlyArray2<'py, f64>,
+        y: PyReadonlyArray1<'py, f64>,
+    ) -> &'py PyArray1<f64> {
+        let x = x.as_array();
+        let y = y.as_array();
+        let z = Zip::from(x.rows()).par_map_collect(|row| row.dot(&y));
+        z.into_pyarray(py)
+    }
+    Ok(())
+}

examples/parallel/tests/test_parallel.py

@@ -0,0 +1,9 @@
+import numpy as np
+import rust_parallel
+
+
+def test_rows_dot():
+    x = np.ones((128, 1024), dtype=np.float64)
+    y = np.ones((1024,), dtype=np.float64)
+    z = rust_parallel.rows_dot(x, y)
+    np.testing.assert_array_almost_equal(z, 1024 * np.ones((128,), dtype=np.float64))

examples/parallel/tox.ini

@@ -0,0 +1,11 @@
+[tox]
+skipsdist = True
+
+[testenv]
+deps =
+    pip
+    numpy
+    pytest
+commands =
+    pip install . -v
+    pytest {posargs}

examples/simple-extension/Cargo.toml

@@ -11,4 +11,3 @@ crate-type = ["cdylib"]
 [dependencies]
 pyo3 = { version = "0.15", features = ["extension-module"] }
 numpy = { path = "../.." }
-num-complex = "0.4.0"

examples/simple-extension/src/lib.rs

@@ -1,6 +1,6 @@
 use numpy::ndarray::{ArrayD, ArrayViewD, ArrayViewMutD};
 use numpy::{Complex64, IntoPyArray, PyArrayDyn, PyReadonlyArrayDyn};
-use pyo3::prelude::{pymodule, PyModule, PyResult, Python};
+use pyo3::{pymodule, types::PyModule, PyResult, Python};
 
 #[pymodule]
 fn rust_ext(_py: Python<'_>, m: &PyModule) -> PyResult<()> {
@@ -30,7 +30,8 @@ fn rust_ext(_py: Python<'_>, m: &PyModule) -> PyResult<()> {
     ) -> &'py PyArrayDyn<f64> {
         let x = x.as_array();
         let y = y.as_array();
-        axpy(a, x, y).into_pyarray(py)
+        let z = axpy(a, x, y);
+        z.into_pyarray(py)
     }
 
     // wrapper of `mult`

examples/simple-extension/tox.ini

@@ -7,5 +7,5 @@ deps =
     numpy
     pytest
 commands =
-    python -m pip install .
+    pip install . -v
     pytest {posargs}

src/array.rs

@@ -244,13 +244,24 @@ impl<T, D> PyArray<T, D> {
         unsafe { Py::from_borrowed_ptr(self.py(), self.as_ptr()) }
     }
 
-    /// Constructs `PyArray` from raw python object without incrementing reference counts.
+    /// Constructs `PyArray` from raw Python object without incrementing reference counts.
+    ///
+    /// # Safety
+    ///
+    /// Implementations must ensure the object does not get freed during `'py`
+    /// and ensure that `ptr` is of the correct type.
     pub unsafe fn from_owned_ptr(py: Python<'_>, ptr: *mut ffi::PyObject) -> &Self {
         py.from_owned_ptr(ptr)
     }
 
-    /// Constructs PyArray from raw python object and increments reference counts.
-    pub unsafe fn from_borrowed_ptr(py: Python<'_>, ptr: *mut ffi::PyObject) -> &Self {
+    /// Constructs PyArray from raw Python object and increments reference counts.
+    ///
+    /// # Safety
+    ///
+    /// Implementations must ensure the object does not get freed during `'py`
+    /// and ensure that `ptr` is of the correct type.
+    /// Note that it must be safe to decrement the reference count of ptr.
+    pub unsafe fn from_borrowed_ptr<'py>(py: Python<'py>, ptr: *mut ffi::PyObject) -> &'py Self {
         py.from_borrowed_ptr(ptr)
     }
 
@@ -673,7 +684,10 @@ impl<T: Element, D: Dimension> PyArray<T, D> {
     ///
     /// See [NpyIndex](../convert/trait.NpyIndex.html) for what types you can use as index.
     ///
-    /// Passing an invalid index can cause undefined behavior(mostly SIGSEGV).
+    /// # Safety
+    ///
+    /// Passing an invalid index is undefined behavior. The element must also have been initialized.
+    /// The elemet must also not be modified by Python code.
     ///
     /// # Example
     /// ```
@@ -693,6 +707,11 @@ impl<T: Element, D: Dimension> PyArray<T, D> {
     }
 
     /// Same as [uget](#method.uget), but returns `&mut T`.
+    ///
+    /// # Safety
+    ///
+    /// Passing an invalid index is undefined behavior. The element must also have been initialized.
+    /// The element must also not be accessed by Python code.
     #[inline(always)]
     #[allow(clippy::mut_from_ref)]
     pub unsafe fn uget_mut<Idx>(&self, index: Idx) -> &mut T
@@ -704,6 +723,9 @@ impl<T: Element, D: Dimension> PyArray<T, D> {
     }
 
     /// Same as [uget](#method.uget), but returns `*mut T`.
+    ///
+    /// # Safety
+    /// Passing an invalid index is undefined behavior.
     #[inline(always)]
     pub unsafe fn uget_raw<Idx>(&self, index: Idx) -> *mut T
     where

src/lib.rs

@@ -1,5 +1,3 @@
-#![allow(clippy::missing_safety_doc)] // FIXME
-
 //! `rust-numpy` provides Rust interfaces for [NumPy C APIs](https://numpy.org/doc/stable/reference/c-api),
 //! especially for [ndarray](https://numpy.org/doc/stable/reference/arrays.ndarray.html) class.
 //!
@@ -29,6 +27,8 @@
 //!     })
 //! }
 //! ```
+#![allow(clippy::needless_lifetimes)] // We often want to make the GIL lifetime explicit.
+
 pub mod array;
 pub mod convert;
 mod dtype;

src/npyffi/mod.rs

@@ -1,7 +1,11 @@
 //! Low-Level bindings for NumPy C API.
 //!
 //! <https://numpy.org/doc/stable/reference/c-api>
-#![allow(non_camel_case_types, clippy::too_many_arguments)]
+#![allow(
+    non_camel_case_types,
+    clippy::too_many_arguments,
+    clippy::missing_safety_doc
+)]
 
 use pyo3::{ffi, Python};
 use std::ffi::CString;

tests/array.rs

@@ -6,11 +6,11 @@ use pyo3::{
     types::{IntoPyDict, PyDict},
 };
 
-fn get_np_locals(py: Python<'_>) -> &'_ PyDict {
+fn get_np_locals(py: Python) -> &PyDict {
     [("np", get_array_module(py).unwrap())].into_py_dict(py)
 }
 
-fn not_contiguous_array<'py>(py: Python<'py>) -> &'py PyArray1<i32> {
+fn not_contiguous_array(py: Python) -> &PyArray1<i32> {
     py.eval(
         "np.array([1, 2, 3, 4], dtype='int32')[::2]",
         Some(get_np_locals(py)),
@@ -266,7 +266,7 @@ fn borrow_from_array() {
     #[pymethods]
     impl Owner {
         #[getter]
-        fn array<'py>(this: &'py PyCell<Self>) -> &'py PyArray1<f64> {
+        fn array(this: &PyCell<Self>) -> &PyArray1<f64> {
             let array = &this.borrow().array;
 
             unsafe { PyArray1::borrow_from_array(array, this) }