diff options
65 files changed, 1516 insertions, 822 deletions
diff --git a/.github/workflows/wheels.yml b/.github/workflows/wheels.yml index 7759e2550..253dd074b 100644 --- a/.github/workflows/wheels.yml +++ b/.github/workflows/wheels.yml @@ -147,3 +147,63 @@ jobs: # https://anaconda.org/multibuild-wheels-staging/numpy # The tokens were originally generated at anaconda.org upload_wheels + build_sdist: + name: Build sdist + needs: get_commit_message + if: >- + contains(needs.get_commit_message.outputs.message, '[wheel build]') || + github.event_name == 'schedule' || + github.event_name == 'workflow_dispatch' || + (github.event_name == 'pull_request' && + (contains(github.event.pull_request.labels.*.name, '36 - Build') || + contains(github.event.pull_request.labels.*.name, '03 - Maintenance') || + contains(github.event.pull_request.labels.*.name, '14 - Release'))) || + (github.event_name == 'push' && startsWith(github.ref, 'refs/tags/v') && ( ! endsWith(github.ref, 'dev0'))) + runs-on: ubuntu-latest + steps: + - name: Checkout numpy + uses: actions/checkout@v2 + with: + submodules: true + # versioneer.py requires the latest tag to be reachable. Here we + # fetch the complete history to get access to the tags. + # A shallow clone can work when the following issue is resolved: + # https://github.com/actions/checkout/issues/338 + fetch-depth: 0 + # Used to push the built wheels + - uses: actions/setup-python@v2 + with: + # Build sdist on lowest supported Python + python-version: '3.8' + - name: Build sdist + run: | + python setup.py sdist + - name: Test the sdist + run: | + # TODO: Don't run test suite, and instead build wheels from sdist + # Depends on pypa/cibuildwheel#1020 + python -m pip install dist/*.gz + pip install -r test_requirements.txt + cd .. # Can't import numpy within numpy src directory + python -c "import numpy; print(numpy.__version__); numpy.test();" + - uses: actions/upload-artifact@v2 + with: + name: sdist + path: ./dist/* + + - name: Upload sdist + if: success() + shell: bash + env: + NUMPY_STAGING_UPLOAD_TOKEN: ${{ secrets.NUMPY_STAGING_UPLOAD_TOKEN }} + NUMPY_NIGHTLY_UPLOAD_TOKEN: ${{ secrets.NUMPY_NIGHTLY_UPLOAD_TOKEN }} + run: | + source tools/wheels/upload_wheels.sh + set_upload_vars + # trigger an upload to + # https://anaconda.org/scipy-wheels-nightly/numpy + # for cron jobs or "Run workflow" (restricted to main branch). + # Tags will upload to + # https://anaconda.org/multibuild-wheels-staging/numpy + # The tokens were originally generated at anaconda.org + upload_wheels @@ -333,6 +333,7 @@ Mattheus Ueckermann <empeeu@yahoo.com> Matthew Barber <quitesimplymatt@gmail.com> Matthew Harrigan <harrigan.matthew@gmail.com> Matthias Bussonnier <bussonniermatthias@gmail.com> <mbussonnier@ucmerced.edu> +Matthieu Darbois <mayeut@users.noreply.github.com> Matti Picus <matti.picus@gmail.com> Maximilian Konrad <maximilianlukaskonrad@hotmail.de> Melissa Weber Mendonça <melissawm@gmail.com> diff --git a/benchmarks/benchmarks/bench_function_base.py b/benchmarks/benchmarks/bench_function_base.py index ac6df784d..d52d7692c 100644 --- a/benchmarks/benchmarks/bench_function_base.py +++ b/benchmarks/benchmarks/bench_function_base.py @@ -284,6 +284,21 @@ class Where(Benchmark): self.d = np.arange(20000) self.e = self.d.copy() self.cond = (self.d > 5000) + size = 1024 * 1024 // 8 + rnd_array = np.random.rand(size) + self.rand_cond_01 = rnd_array > 0.01 + self.rand_cond_20 = rnd_array > 0.20 + self.rand_cond_30 = rnd_array > 0.30 + self.rand_cond_40 = rnd_array > 0.40 + self.rand_cond_50 = rnd_array > 0.50 + self.all_zeros = np.zeros(size, dtype=bool) + self.all_ones = np.ones(size, dtype=bool) + self.rep_zeros_2 = np.arange(size) % 2 == 0 + self.rep_zeros_4 = np.arange(size) % 4 == 0 + self.rep_zeros_8 = np.arange(size) % 8 == 0 + self.rep_ones_2 = np.arange(size) % 2 > 0 + self.rep_ones_4 = np.arange(size) % 4 > 0 + self.rep_ones_8 = np.arange(size) % 8 > 0 def time_1(self): np.where(self.cond) @@ -293,3 +308,43 @@ class Where(Benchmark): def time_2_broadcast(self): np.where(self.cond, self.d, 0) + + def time_all_zeros(self): + np.where(self.all_zeros) + + def time_random_01_percent(self): + np.where(self.rand_cond_01) + + def time_random_20_percent(self): + np.where(self.rand_cond_20) + + def time_random_30_percent(self): + np.where(self.rand_cond_30) + + def time_random_40_percent(self): + np.where(self.rand_cond_40) + + def time_random_50_percent(self): + np.where(self.rand_cond_50) + + def time_all_ones(self): + np.where(self.all_ones) + + def time_interleaved_zeros_x2(self): + np.where(self.rep_zeros_2) + + def time_interleaved_zeros_x4(self): + np.where(self.rep_zeros_4) + + def time_interleaved_zeros_x8(self): + np.where(self.rep_zeros_8) + + def time_interleaved_ones_x2(self): + np.where(self.rep_ones_2) + + def time_interleaved_ones_x4(self): + np.where(self.rep_ones_4) + + def time_interleaved_ones_x8(self): + np.where(self.rep_ones_8) + diff --git a/doc/changelog/1.22.3-changelog.rst b/doc/changelog/1.22.3-changelog.rst new file mode 100644 index 000000000..051e5665b --- /dev/null +++ b/doc/changelog/1.22.3-changelog.rst @@ -0,0 +1,32 @@ + +Contributors +============ + +A total of 9 people contributed to this release. People with a "+" by their +names contributed a patch for the first time. + +* @GalaxySnail + +* Alexandre de Siqueira +* Bas van Beek +* Charles Harris +* Melissa Weber Mendonça +* Ross Barnowski +* Sebastian Berg +* Tirth Patel +* Matthieu Darbois + +Pull requests merged +==================== + +A total of 10 pull requests were merged for this release. + +* `#21048 <https://github.com/numpy/numpy/pull/21048>`__: MAINT: Use "3.10" instead of "3.10-dev" on travis. +* `#21106 <https://github.com/numpy/numpy/pull/21106>`__: TYP,MAINT: Explicitly allow sequences of array-likes in ``np.concatenate`` +* `#21137 <https://github.com/numpy/numpy/pull/21137>`__: BLD,DOC: skip broken ipython 8.1.0 +* `#21138 <https://github.com/numpy/numpy/pull/21138>`__: BUG, ENH: np._from_dlpack: export correct device information +* `#21139 <https://github.com/numpy/numpy/pull/21139>`__: BUG: Fix numba DUFuncs added loops getting picked up +* `#21140 <https://github.com/numpy/numpy/pull/21140>`__: BUG: Fix unpickling an empty ndarray with a non-zero dimension... +* `#21141 <https://github.com/numpy/numpy/pull/21141>`__: BUG: use ThreadPoolExecutor instead of ThreadPool +* `#21142 <https://github.com/numpy/numpy/pull/21142>`__: API: Disallow strings in logical ufuncs +* `#21143 <https://github.com/numpy/numpy/pull/21143>`__: MAINT, DOC: Fix SciPy intersphinx link +* `#21148 <https://github.com/numpy/numpy/pull/21148>`__: BUG,ENH: np._from_dlpack: export arrays with any strided size-1... diff --git a/doc/neps/nep-0029-deprecation_policy.rst b/doc/neps/nep-0029-deprecation_policy.rst index a50afcb98..36f128159 100644 --- a/doc/neps/nep-0029-deprecation_policy.rst +++ b/doc/neps/nep-0029-deprecation_policy.rst @@ -114,7 +114,12 @@ Jul 26, 2021 3.7+ 1.18+ Dec 22, 2021 3.7+ 1.19+ Dec 26, 2021 3.8+ 1.19+ Jun 21, 2022 3.8+ 1.20+ -Apr 14, 2023 3.9+ 1.20+ +Jan 31, 2023 3.8+ 1.21+ +Apr 14, 2023 3.9+ 1.21+ +Jun 23, 2023 3.9+ 1.22+ +Jan 01, 2024 3.9+ 1.23+ +Apr 05, 2024 3.10+ 1.23+ +Apr 04, 2025 3.11+ 1.23+ ============ ====== ===== @@ -132,7 +137,12 @@ Drop Schedule On Dec 22, 2021 drop support for NumPy 1.18 (initially released on Dec 22, 2019) On Dec 26, 2021 drop support for Python 3.7 (initially released on Jun 27, 2018) On Jun 21, 2022 drop support for NumPy 1.19 (initially released on Jun 20, 2020) + On Jan 31, 2023 drop support for NumPy 1.20 (initially released on Jan 31, 2021) On Apr 14, 2023 drop support for Python 3.8 (initially released on Oct 14, 2019) + On Jun 23, 2023 drop support for NumPy 1.21 (initially released on Jun 22, 2021) + On Jan 01, 2024 drop support for NumPy 1.22 (initially released on Dec 31, 2021) + On Apr 05, 2024 drop support for Python 3.9 (initially released on Oct 05, 2020) + On Apr 04, 2025 drop support for Python 3.10 (initially released on Oct 04, 2021) Implementation @@ -261,6 +271,11 @@ Code to generate support and drop schedule tables :: Oct 14, 2019: Python 3.8 Dec 22, 2019: NumPy 1.18 Jun 20, 2020: NumPy 1.19 + Oct 05, 2020: Python 3.9 + Jan 30, 2021: NumPy 1.20 + Jun 22, 2021: NumPy 1.21 + Oct 04, 2021: Python 3.10 + Dec 31, 2021: NumPy 1.22 """ releases = [] diff --git a/doc/neps/nep-0047-array-api-standard.rst b/doc/neps/nep-0047-array-api-standard.rst index 53b8e35b0..a94b3b423 100644 --- a/doc/neps/nep-0047-array-api-standard.rst +++ b/doc/neps/nep-0047-array-api-standard.rst @@ -340,7 +340,7 @@ Adding support for DLPack to NumPy entails: - Adding a ``ndarray.__dlpack__()`` method which returns a ``dlpack`` C structure wrapped in a ``PyCapsule``. -- Adding a ``np._from_dlpack(obj)`` function, where ``obj`` supports +- Adding a ``np.from_dlpack(obj)`` function, where ``obj`` supports ``__dlpack__()``, and returns an ``ndarray``. DLPack is currently a ~200 LoC header, and is meant to be included directly, so diff --git a/doc/release/upcoming_changes/21001.performance.rst b/doc/release/upcoming_changes/21001.performance.rst new file mode 100644 index 000000000..04dc43482 --- /dev/null +++ b/doc/release/upcoming_changes/21001.performance.rst @@ -0,0 +1,5 @@ +Faster reduction operators +-------------------------- +Reduction operations like `numpy.sum`, `numpy.prod`, `numpy.add.reduce`, +`numpy.logical_and.reduce` on contiguous integer-based arrays are now +much faster. diff --git a/doc/release/upcoming_changes/21130.performance.rst b/doc/release/upcoming_changes/21130.performance.rst new file mode 100644 index 000000000..70ea7dca4 --- /dev/null +++ b/doc/release/upcoming_changes/21130.performance.rst @@ -0,0 +1,4 @@ +Faster ``np.where`` +------------------- +`numpy.where` is now much faster than previously on unpredictable/random +input data. diff --git a/doc/release/upcoming_changes/21145.new_function.rst b/doc/release/upcoming_changes/21145.new_function.rst new file mode 100644 index 000000000..75fa9e181 --- /dev/null +++ b/doc/release/upcoming_changes/21145.new_function.rst @@ -0,0 +1,6 @@ +NumPy now supports the DLPack protocol +-------------------------------------- +`numpy.from_dlpack` has been added to NumPy to exchange data using the DLPack protocol. +It accepts Python objects that implement the ``__dlpack__`` and ``__dlpack_device__`` +methods and returns a ndarray object which is generally the view of the data of the input +object. diff --git a/doc/source/conf.py b/doc/source/conf.py index 5b7c2a5e6..4301fe553 100644 --- a/doc/source/conf.py +++ b/doc/source/conf.py @@ -307,6 +307,7 @@ intersphinx_mapping = { 'pytest': ('https://docs.pytest.org/en/stable', None), 'numpy-tutorials': ('https://numpy.org/numpy-tutorials', None), 'numpydoc': ('https://numpydoc.readthedocs.io/en/latest', None), + 'dlpack': ('https://dmlc.github.io/dlpack/latest', None) } diff --git a/doc/source/dev/development_workflow.rst b/doc/source/dev/development_workflow.rst index 0e46341b9..502c89939 100644 --- a/doc/source/dev/development_workflow.rst +++ b/doc/source/dev/development_workflow.rst @@ -218,12 +218,27 @@ cibuildwheel wheel builders are not run by default on every single PR or commit If you would like to test that your pull request do not break the wheel builders, you may either append ``[wheel build]`` to the end of the commit message of the commit -or add one of the following labels to the pull request(if you have the permissions to do so):: +or add one of the following labels to the pull request(if you have the permissions to do so): - ``36 - Build``: for pull requests changing build processes/configurations - ``03 - Maintenance``: for pull requests upgrading dependencies - ``14 - Release``: for pull requests preparing for a release +- ``36 - Build``: for pull requests changing build processes/configurations +- ``03 - Maintenance``: for pull requests upgrading dependencies +- ``14 - Release``: for pull requests preparing for a release +The wheels built via github actions (including 64-bit linux, macOS, and +windows, arm64 macOS, and 32-bit windows) will be uploaded as artifacts in zip +files. You can access them from the Summary page of the "Wheel builder" +Action_. The aarch64 wheels built via travis_ CI are not available as artifacts. +Additionally, the wheels will be uploaded to +https://anaconda.org/scipy-wheels-nightly/ on the following conditions: + +- by a weekly cron job or +- if the github action or travis build has been manually triggered, which requires appropriate permissions + +The wheels wil be uploaded to https://anaconda.org/multibuild-wheels-staging/ +if the build was triggered by a tag to the repo that begins with ``v`` + +.. _Action: https://github.com/numpy/numpy/actions +.. _travis: https://app.travis-ci.com/github/numpy/numpy/builds .. _workflow_mailing_list: diff --git a/doc/source/reference/arrays.interface.rst b/doc/source/reference/arrays.interface.rst index 25cfa3de1..e10710719 100644 --- a/doc/source/reference/arrays.interface.rst +++ b/doc/source/reference/arrays.interface.rst @@ -91,7 +91,7 @@ This approach to the interface consists of the object having an ``M`` Datetime ``O`` Object (i.e. the memory contains a pointer to :c:type:`PyObject`) ``S`` String (fixed-length sequence of char) - ``U`` Unicode (fixed-length sequence of :c:type:`Py_UNICODE`) + ``U`` Unicode (fixed-length sequence of :c:type:`Py_UCS4`) ``V`` Other (void \* -- each item is a fixed-size chunk of memory) ===== ================================================================ @@ -247,7 +247,8 @@ flag is present. .. note:: :obj:`__array_struct__` is considered legacy and should not be used for new - code. Use the :py:doc:`buffer protocol <c-api/buffer>` instead. + code. Use the :py:doc:`buffer protocol <c-api/buffer>` or the DLPack protocol + `numpy.from_dlpack` instead. Type description examples diff --git a/doc/source/reference/routines.array-creation.rst b/doc/source/reference/routines.array-creation.rst index 30780c286..9d2954f2c 100644 --- a/doc/source/reference/routines.array-creation.rst +++ b/doc/source/reference/routines.array-creation.rst @@ -35,6 +35,7 @@ From existing data asmatrix copy frombuffer + from_dlpack fromfile fromfunction fromiter diff --git a/doc/source/release.rst b/doc/source/release.rst index 42bc2fb07..3fb8ad287 100644 --- a/doc/source/release.rst +++ b/doc/source/release.rst @@ -6,6 +6,7 @@ Release notes :maxdepth: 3 1.23.0 <release/1.23.0-notes> + 1.22.3 <release/1.22.3-notes> 1.22.2 <release/1.22.2-notes> 1.22.1 <release/1.22.1-notes> 1.22.0 <release/1.22.0-notes> diff --git a/doc/source/release/1.22.0-notes.rst b/doc/source/release/1.22.0-notes.rst index 2cbb24b60..c5477fb16 100644 --- a/doc/source/release/1.22.0-notes.rst +++ b/doc/source/release/1.22.0-notes.rst @@ -27,11 +27,13 @@ These are in addition to the ongoing work to provide SIMD support for commonly used functions, improvements to F2PY, and better documentation. The Python versions supported in this release are 3.8-3.10, Python 3.7 has been -dropped. Note that 32 bit wheels are only provided for Python 3.8 and 3.9 on -Windows, all other wheels are 64 bits on account of Ubuntu, Fedora, and other -Linux distributions dropping 32 bit support. All 64 bit wheels are also linked -with 64 bit integer OpenBLAS, which should fix the occasional problems -encountered by folks using truly huge arrays. +dropped. Note that the Mac wheels are now based on OS X 10.14 rather than 10.9 +that was used in previous NumPy release cycles. 10.14 is the oldest release +supported by Apple. Also note that 32 bit wheels are only provided for Python +3.8 and 3.9 on Windows, all other wheels are 64 bits on account of Ubuntu, +Fedora, and other Linux distributions dropping 32 bit support. All 64 bit +wheels are also linked with 64 bit integer OpenBLAS, which should fix the +occasional problems encountered by folks using truly huge arrays. Expired deprecations diff --git a/doc/source/release/1.22.3-notes.rst b/doc/source/release/1.22.3-notes.rst new file mode 100644 index 000000000..6c94e2d28 --- /dev/null +++ b/doc/source/release/1.22.3-notes.rst @@ -0,0 +1,48 @@ +.. currentmodule:: numpy + +========================== +NumPy 1.22.3 Release Notes +========================== + +NumPy 1.22.3 is a maintenance release that fixes bugs discovered after the +1.22.2 release. The most noticeable fixes may be those for DLPack. One that may +cause some problems is disallowing strings as inputs to logical ufuncs. It is +still undecided how strings should be treated in those functions and it was +thought best to simply disallow them until a decision was reached. That should +not cause problems with older code. + +The Python versions supported for this release are 3.8-3.10. Note that the Mac +wheels are now based on OS X 10.14 rather than 10.9 that was used in previous +NumPy release cycles. 10.14 is the oldest release supported by Apple. + +Contributors +============ + +A total of 9 people contributed to this release. People with a "+" by their +names contributed a patch for the first time. + +* @GalaxySnail + +* Alexandre de Siqueira +* Bas van Beek +* Charles Harris +* Melissa Weber Mendonça +* Ross Barnowski +* Sebastian Berg +* Tirth Patel +* Matthieu Darbois + +Pull requests merged +==================== + +A total of 10 pull requests were merged for this release. + +* `#21048 <https://github.com/numpy/numpy/pull/21048>`__: MAINT: Use "3.10" instead of "3.10-dev" on travis. +* `#21106 <https://github.com/numpy/numpy/pull/21106>`__: TYP,MAINT: Explicitly allow sequences of array-likes in ``np.concatenate`` +* `#21137 <https://github.com/numpy/numpy/pull/21137>`__: BLD,DOC: skip broken ipython 8.1.0 +* `#21138 <https://github.com/numpy/numpy/pull/21138>`__: BUG, ENH: np._from_dlpack: export correct device information +* `#21139 <https://github.com/numpy/numpy/pull/21139>`__: BUG: Fix numba DUFuncs added loops getting picked up +* `#21140 <https://github.com/numpy/numpy/pull/21140>`__: BUG: Fix unpickling an empty ndarray with a non-zero dimension... +* `#21141 <https://github.com/numpy/numpy/pull/21141>`__: BUG: use ThreadPoolExecutor instead of ThreadPool +* `#21142 <https://github.com/numpy/numpy/pull/21142>`__: API: Disallow strings in logical ufuncs +* `#21143 <https://github.com/numpy/numpy/pull/21143>`__: MAINT, DOC: Fix SciPy intersphinx link +* `#21148 <https://github.com/numpy/numpy/pull/21148>`__: BUG,ENH: np._from_dlpack: export arrays with any strided size-1... diff --git a/doc/source/user/basics.interoperability.rst b/doc/source/user/basics.interoperability.rst index adad4dab9..853f324ba 100644 --- a/doc/source/user/basics.interoperability.rst +++ b/doc/source/user/basics.interoperability.rst @@ -55,6 +55,14 @@ describes its memory layout and NumPy does everything else (zero-copy if possible). If that's not possible, the object itself is responsible for returning a ``ndarray`` from ``__array__()``. +:doc:`DLPack <dlpack:index>` is yet another protocol to convert foriegn objects +to NumPy arrays in a language and device agnostic manner. NumPy doesn't implicitly +convert objects to ndarrays using DLPack. It provides the function +`numpy.from_dlpack` that accepts any object implementing the ``__dlpack__`` method +and outputs a NumPy ndarray (which is generally a view of the input object's data +buffer). The :ref:`dlpack:python-spec` page explains the ``__dlpack__`` protocol +in detail. + The array interface protocol ~~~~~~~~~~~~~~~~~~~~~~~~~~~~ @@ -118,6 +126,26 @@ as the original object and any attributes/behavior it may have had, is lost. To see an example of a custom array implementation including the use of ``__array__()``, see :ref:`basics.dispatch`. +The DLPack Protocol +~~~~~~~~~~~~~~~~~~~ + +The :doc:`DLPack <dlpack:index>` protocol defines a memory-layout of +strided n-dimensional array objects. It offers the following syntax +for data exchange: + +1. A `numpy.from_dlpack` function, which accepts (array) objects with a + ``__dlpack__`` method and uses that method to construct a new array + containing the data from ``x``. +2. ``__dlpack__(self, stream=None)`` and ``__dlpack_device__`` methods on the + array object, which will be called from within ``from_dlpack``, to query + what device the array is on (may be needed to pass in the correct + stream, e.g. in the case of multiple GPUs) and to access the data. + +Unlike the buffer protocol, DLPack allows exchanging arrays containing data on +devices other than the CPU (e.g. Vulkan or GPU). Since NumPy only supports CPU, +it can only convert objects whose data exists on the CPU. But other libraries, +like PyTorch_ and CuPy_, may exchange data on GPU using this protocol. + 2. Operating on foreign objects without converting -------------------------------------------------- @@ -395,6 +423,78 @@ See `the Dask array documentation and the `scope of Dask arrays interoperability with NumPy arrays <https://docs.dask.org/en/stable/array.html#scope>`__ for details. +Example: DLPack +~~~~~~~~~~~~~~~ + +Several Python data science libraries implement the ``__dlpack__`` protocol. +Among them are PyTorch_ and CuPy_. A full list of libraries that implement +this protocol can be found on +:doc:`this page of DLPack documentation <dlpack:index>`. + +Convert a PyTorch CPU tensor to NumPy array: + + >>> import torch + >>> x_torch = torch.arange(5) + >>> x_torch + tensor([0, 1, 2, 3, 4]) + >>> x_np = np.from_dlpack(x_torch) + >>> x_np + array([0, 1, 2, 3, 4]) + >>> # note that x_np is a view of x_torch + >>> x_torch[1] = 100 + >>> x_torch + tensor([ 0, 100, 2, 3, 4]) + >>> x_np + array([ 0, 100, 2, 3, 4]) + +The imported arrays are read-only so writing or operating in-place will fail: + + >>> x.flags.writeable + False + >>> x_np[1] = 1 + Traceback (most recent call last): + File "<stdin>", line 1, in <module> + ValueError: assignment destination is read-only + +A copy must be created in order to operate on the imported arrays in-place, but +will mean duplicating the memory. Do not do this for very large arrays: + + >>> x_np_copy = x_np.copy() + >>> x_np_copy.sort() # works + +.. note:: + + Note that GPU tensors can't be converted to NumPy arrays since NumPy doesn't + support GPU devices: + + >>> x_torch = torch.arange(5, device='cuda') + >>> np.from_dlpack(x_torch) + Traceback (most recent call last): + File "<stdin>", line 1, in <module> + RuntimeError: Unsupported device in DLTensor. + + But, if both libraries support the device the data buffer is on, it is + possible to use the ``__dlpack__`` protocol (e.g. PyTorch_ and CuPy_): + + >>> x_torch = torch.arange(5, device='cuda') + >>> x_cupy = cupy.from_dlpack(x_torch) + +Similarly, a NumPy array can be converted to a PyTorch tensor: + + >>> x_np = np.arange(5) + >>> x_torch = torch.from_dlpack(x_np) + +Read-only arrays cannot be exported: + + >>> x_np = np.arange(5) + >>> x_np.flags.writeable = False + >>> torch.from_dlpack(x_np) # doctest: +ELLIPSIS + Traceback (most recent call last): + File "<stdin>", line 1, in <module> + File ".../site-packages/torch/utils/dlpack.py", line 63, in from_dlpack + dlpack = ext_tensor.__dlpack__() + TypeError: NumPy currently only supports dlpack for writeable arrays + Further reading --------------- diff --git a/doc/source/user/building.rst b/doc/source/user/building.rst index 82fc1911f..94ae1ea12 100644 --- a/doc/source/user/building.rst +++ b/doc/source/user/building.rst @@ -8,7 +8,7 @@ source. Your choice depends on your operating system and familiarity with the command line. Gitpod ------------- +------ Gitpod is an open-source platform that automatically creates the correct development environment right in your browser, reducing the need to @@ -21,7 +21,7 @@ in-depth instructions for building NumPy with `building NumPy with Gitpod`_. .. _building NumPy with Gitpod: https://numpy.org/devdocs/dev/development_gitpod.html Building locally ------------------- +---------------- Building locally on your machine gives you more granular control. If you are a MacOS or Linux user familiar with using the @@ -94,7 +94,14 @@ Testing ------- Make sure to test your builds. To ensure everything stays in shape, see if -all tests pass:: +all tests pass. + +The test suite requires additional dependencies, which can easily be +installed with:: + + $ python -m pip install -r test_requirements.txt + +Run tests:: $ python runtests.py -v -m full diff --git a/doc/source/user/depending_on_numpy.rst b/doc/source/user/depending_on_numpy.rst index d8e97ef1f..c61b1d7fe 100644 --- a/doc/source/user/depending_on_numpy.rst +++ b/doc/source/user/depending_on_numpy.rst @@ -114,26 +114,30 @@ for dropping support for old Python and NumPy versions: :ref:`NEP29`. We recommend all packages depending on NumPy to follow the recommendations in NEP 29. -For *run-time dependencies*, you specify the range of versions in +For *run-time dependencies*, specify version bounds using ``install_requires`` in ``setup.py`` (assuming you use ``numpy.distutils`` or -``setuptools`` to build). Getting the upper bound right for NumPy is slightly -tricky. If we don't set any bound, a too-new version will be pulled in a few -years down the line, and NumPy may have deprecated and removed some API that -your package depended on by then. On the other hand if you set the upper bound -to the newest already-released version, then as soon as a new NumPy version is -released there will be no matching version of your package that works with it. - -What to do here depends on your release frequency. Given that NumPy releases -come in a 6-monthly cadence and that features that get deprecated in NumPy -should stay around for another two releases, a good upper bound is -``<1.(xx+3).0`` - where ``xx`` is the minor version of the latest -already-released NumPy. This is safe to do if you release at least once a year. -If your own releases are much less frequent, you may set the upper bound a -little further into the future - this is a trade-off between a future NumPy -version _maybe_ removing something you rely on, and the upper bound being -exceeded which _may_ lead to your package being hard to install in combination -with other packages relying on the latest NumPy. - +``setuptools`` to build). + +Most libraries that rely on NumPy will not need to set an upper +version bound: NumPy is careful to preserve backward-compatibility. + +That said, if you are (a) a project that is guaranteed to release +frequently, (b) use a large part of NumPy's API surface, and (c) is +worried that changes in NumPy may break your code, you can set an +upper bound of ``<MAJOR.MINOR + N`` with N no less than 3, and +``MAJOR.MINOR`` being the current release of NumPy [*]_. If you use the NumPy +C API (directly or via Cython), you can also pin the current major +version to prevent ABI breakage. Note that setting an upper bound on +NumPy may `affect the ability of your library to be installed +alongside other, newer packages +<https://iscinumpy.dev/post/bound-version-constraints/>`__. + +.. [*] The reason for setting ``N=3`` is that NumPy will, on the + rare occasion where it makes breaking changes, raise warnings + for at least two releases. (NumPy releases about once every six + months, so this translates to a window of at least a year; + hence the subsequent requirement that your project releases at + least on that cadence.) .. note:: diff --git a/environment.yml b/environment.yml index 7d8fba983..db1b42ee1 100644 --- a/environment.yml +++ b/environment.yml @@ -19,7 +19,7 @@ dependencies: - pytest-xdist - hypothesis # For type annotations - - mypy=0.931 + - mypy=0.940 # For building docs - sphinx=4.2.0 - sphinx-panels diff --git a/numpy/__init__.pyi b/numpy/__init__.pyi index 297c482e5..59a0c6673 100644 --- a/numpy/__init__.pyi +++ b/numpy/__init__.pyi @@ -4372,4 +4372,4 @@ class chararray(ndarray[_ShapeType, _CharDType]): class _SupportsDLPack(Protocol[_T_contra]): def __dlpack__(self, *, stream: None | _T_contra = ...) -> _PyCapsule: ... -def _from_dlpack(__obj: _SupportsDLPack[None]) -> NDArray[Any]: ... +def from_dlpack(__obj: _SupportsDLPack[None]) -> NDArray[Any]: ... diff --git a/numpy/array_api/_creation_functions.py b/numpy/array_api/_creation_functions.py index 741498ff6..3b014d37b 100644 --- a/numpy/array_api/_creation_functions.py +++ b/numpy/array_api/_creation_functions.py @@ -154,7 +154,7 @@ def eye( def from_dlpack(x: object, /) -> Array: from ._array_object import Array - return Array._new(np._from_dlpack(x)) + return Array._new(np.from_dlpack(x)) def full( diff --git a/numpy/core/_add_newdocs.py b/numpy/core/_add_newdocs.py index 6ac9951fb..80409669d 100644 --- a/numpy/core/_add_newdocs.py +++ b/numpy/core/_add_newdocs.py @@ -1573,17 +1573,38 @@ add_newdoc('numpy.core.multiarray', 'frombuffer', array_function_like_doc, )) -add_newdoc('numpy.core.multiarray', '_from_dlpack', +add_newdoc('numpy.core.multiarray', 'from_dlpack', """ - _from_dlpack(x, /) + from_dlpack(x, /) Create a NumPy array from an object implementing the ``__dlpack__`` - protocol. + protocol. Generally, the returned NumPy array is a read-only view + of the input object. See [1]_ and [2]_ for more details. - See Also + Parameters + ---------- + x : object + A Python object that implements the ``__dlpack__`` and + ``__dlpack_device__`` methods. + + Returns + ------- + out : ndarray + + References + ---------- + .. [1] Array API documentation, + https://data-apis.org/array-api/latest/design_topics/data_interchange.html#syntax-for-data-interchange-with-dlpack + + .. [2] Python specification for DLPack, + https://dmlc.github.io/dlpack/latest/python_spec.html + + Examples -------- - `Array API documentation - <https://data-apis.org/array-api/latest/design_topics/data_interchange.html#syntax-for-data-interchange-with-dlpack>`_ + >>> import torch + >>> x = torch.arange(10) + >>> # create a view of the torch tensor "x" in NumPy + >>> y = np.from_dlpack(x) """) add_newdoc('numpy.core', 'fastCopyAndTranspose', diff --git a/numpy/core/_type_aliases.py b/numpy/core/_type_aliases.py index 3765a0d34..9b1dcb8cf 100644 --- a/numpy/core/_type_aliases.py +++ b/numpy/core/_type_aliases.py @@ -172,7 +172,7 @@ def _set_up_aliases(): allTypes[alias] = allTypes[t] sctypeDict[alias] = sctypeDict[t] # Remove aliases overriding python types and modules - to_remove = ['ulong', 'object', 'int', 'float', + to_remove = ['object', 'int', 'float', 'complex', 'bool', 'string', 'datetime', 'timedelta', 'bytes', 'str'] @@ -182,6 +182,15 @@ def _set_up_aliases(): del sctypeDict[t] except KeyError: pass + + # Additional aliases in sctypeDict that should not be exposed as attributes + attrs_to_remove = ['ulong'] + + for t in attrs_to_remove: + try: + del allTypes[t] + except KeyError: + pass _set_up_aliases() diff --git a/numpy/core/fromnumeric.py b/numpy/core/fromnumeric.py index 0ec4f7687..9e58d9bea 100644 --- a/numpy/core/fromnumeric.py +++ b/numpy/core/fromnumeric.py @@ -804,8 +804,8 @@ def argpartition(a, kth, axis=-1, kind='introselect', order=None): index_array : ndarray, int Array of indices that partition `a` along the specified axis. If `a` is one-dimensional, ``a[index_array]`` yields a partitioned `a`. - More generally, ``np.take_along_axis(a, index_array, axis=a)`` always - yields the partitioned `a`, irrespective of dimensionality. + More generally, ``np.take_along_axis(a, index_array, axis)`` + always yields the partitioned `a`, irrespective of dimensionality. See Also -------- @@ -2309,7 +2309,7 @@ def any(a, axis=None, out=None, keepdims=np._NoValue, *, where=np._NoValue): """ Test whether any array element along a given axis evaluates to True. - Returns single boolean unless `axis` is not ``None`` + Returns single boolean if `axis` is ``None`` Parameters ---------- diff --git a/numpy/core/multiarray.py b/numpy/core/multiarray.py index f88d75978..1a37ed3e1 100644 --- a/numpy/core/multiarray.py +++ b/numpy/core/multiarray.py @@ -14,7 +14,7 @@ from ._multiarray_umath import * # noqa: F403 # do not change them. issue gh-15518 # _get_ndarray_c_version is semi-public, on purpose not added to __all__ from ._multiarray_umath import ( - _fastCopyAndTranspose, _flagdict, _from_dlpack, _insert, _reconstruct, + _fastCopyAndTranspose, _flagdict, from_dlpack, _insert, _reconstruct, _vec_string, _ARRAY_API, _monotonicity, _get_ndarray_c_version, _set_madvise_hugepage, ) @@ -24,7 +24,7 @@ __all__ = [ 'ITEM_HASOBJECT', 'ITEM_IS_POINTER', 'LIST_PICKLE', 'MAXDIMS', 'MAY_SHARE_BOUNDS', 'MAY_SHARE_EXACT', 'NEEDS_INIT', 'NEEDS_PYAPI', 'RAISE', 'USE_GETITEM', 'USE_SETITEM', 'WRAP', '_fastCopyAndTranspose', - '_flagdict', '_from_dlpack', '_insert', '_reconstruct', '_vec_string', + '_flagdict', 'from_dlpack', '_insert', '_reconstruct', '_vec_string', '_monotonicity', 'add_docstring', 'arange', 'array', 'asarray', 'asanyarray', 'ascontiguousarray', 'asfortranarray', 'bincount', 'broadcast', 'busday_count', 'busday_offset', 'busdaycalendar', 'can_cast', @@ -47,7 +47,7 @@ _reconstruct.__module__ = 'numpy.core.multiarray' scalar.__module__ = 'numpy.core.multiarray' -_from_dlpack.__module__ = 'numpy' +from_dlpack.__module__ = 'numpy' arange.__module__ = 'numpy' array.__module__ = 'numpy' asarray.__module__ = 'numpy' diff --git a/numpy/core/numeric.py b/numpy/core/numeric.py index 5f6602c71..39bfc2d55 100644 --- a/numpy/core/numeric.py +++ b/numpy/core/numeric.py @@ -13,7 +13,7 @@ from .multiarray import ( WRAP, arange, array, asarray, asanyarray, ascontiguousarray, asfortranarray, broadcast, can_cast, compare_chararrays, concatenate, copyto, dot, dtype, empty, - empty_like, flatiter, frombuffer, _from_dlpack, fromfile, fromiter, + empty_like, flatiter, frombuffer, from_dlpack, fromfile, fromiter, fromstring, inner, lexsort, matmul, may_share_memory, min_scalar_type, ndarray, nditer, nested_iters, promote_types, putmask, result_type, set_numeric_ops, shares_memory, vdot, where, @@ -41,7 +41,7 @@ __all__ = [ 'newaxis', 'ndarray', 'flatiter', 'nditer', 'nested_iters', 'ufunc', 'arange', 'array', 'asarray', 'asanyarray', 'ascontiguousarray', 'asfortranarray', 'zeros', 'count_nonzero', 'empty', 'broadcast', 'dtype', - 'fromstring', 'fromfile', 'frombuffer', '_from_dlpack', 'where', + 'fromstring', 'fromfile', 'frombuffer', 'from_dlpack', 'where', 'argwhere', 'copyto', 'concatenate', 'fastCopyAndTranspose', 'lexsort', 'set_numeric_ops', 'can_cast', 'promote_types', 'min_scalar_type', 'result_type', 'isfortran', 'empty_like', 'zeros_like', 'ones_like', diff --git a/numpy/core/setup.py b/numpy/core/setup.py index f6b31075d..c4222d7c0 100644 --- a/numpy/core/setup.py +++ b/numpy/core/setup.py @@ -947,7 +947,7 @@ def configuration(parent_package='',top_path=None): join('src', 'multiarray', 'usertypes.c'), join('src', 'multiarray', 'vdot.c'), join('src', 'common', 'npy_sort.h.src'), - join('src', 'npysort', 'x86-qsort.dispatch.c.src'), + join('src', 'npysort', 'x86-qsort.dispatch.cpp'), join('src', 'npysort', 'quicksort.cpp'), join('src', 'npysort', 'mergesort.cpp'), join('src', 'npysort', 'timsort.cpp'), diff --git a/numpy/core/src/common/npy_dlpack.h b/numpy/core/src/common/npy_dlpack.h index 14ca352c0..cb926a262 100644 --- a/numpy/core/src/common/npy_dlpack.h +++ b/numpy/core/src/common/npy_dlpack.h @@ -23,6 +23,6 @@ array_dlpack_device(PyArrayObject *self, PyObject *NPY_UNUSED(args)); NPY_NO_EXPORT PyObject * -_from_dlpack(PyObject *NPY_UNUSED(self), PyObject *obj); +from_dlpack(PyObject *NPY_UNUSED(self), PyObject *obj); #endif diff --git a/numpy/core/src/common/simd/avx512/memory.h b/numpy/core/src/common/simd/avx512/memory.h index dcfb6c890..03fcb4630 100644 --- a/numpy/core/src/common/simd/avx512/memory.h +++ b/numpy/core/src/common/simd/avx512/memory.h @@ -276,7 +276,8 @@ NPY_FINLINE void npyv_storen_till_s64(npy_int64 *ptr, npy_intp stride, npy_uintp union { \ npyv_lanetype_##F_SFX from_##F_SFX; \ npyv_lanetype_##T_SFX to_##T_SFX; \ - } pun = {.from_##F_SFX = fill}; \ + } pun; \ + pun.from_##F_SFX = fill; \ return npyv_reinterpret_##F_SFX##_##T_SFX(npyv_load_till_##T_SFX( \ (const npyv_lanetype_##T_SFX *)ptr, nlane, pun.to_##T_SFX \ )); \ @@ -288,7 +289,8 @@ NPY_FINLINE void npyv_storen_till_s64(npy_int64 *ptr, npy_intp stride, npy_uintp union { \ npyv_lanetype_##F_SFX from_##F_SFX; \ npyv_lanetype_##T_SFX to_##T_SFX; \ - } pun = {.from_##F_SFX = fill}; \ + } pun; \ + pun.from_##F_SFX = fill; \ return npyv_reinterpret_##F_SFX##_##T_SFX(npyv_loadn_till_##T_SFX( \ (const npyv_lanetype_##T_SFX *)ptr, stride, nlane, pun.to_##T_SFX \ )); \ diff --git a/numpy/core/src/common/simd/avx512/reorder.h b/numpy/core/src/common/simd/avx512/reorder.h index f043004ec..c0b2477f3 100644 --- a/numpy/core/src/common/simd/avx512/reorder.h +++ b/numpy/core/src/common/simd/avx512/reorder.h @@ -214,13 +214,13 @@ NPY_FINLINE npyv_u16 npyv_rev64_u16(npyv_u16 a) NPY_FINLINE npyv_u32 npyv_rev64_u32(npyv_u32 a) { - return _mm512_shuffle_epi32(a, _MM_SHUFFLE(2, 3, 0, 1)); + return _mm512_shuffle_epi32(a, (_MM_PERM_ENUM)_MM_SHUFFLE(2, 3, 0, 1)); } #define npyv_rev64_s32 npyv_rev64_u32 NPY_FINLINE npyv_f32 npyv_rev64_f32(npyv_f32 a) { - return _mm512_shuffle_ps(a, a, _MM_SHUFFLE(2, 3, 0, 1)); + return _mm512_shuffle_ps(a, a, (_MM_PERM_ENUM)_MM_SHUFFLE(2, 3, 0, 1)); } #endif // _NPY_SIMD_AVX512_REORDER_H diff --git a/numpy/core/src/multiarray/dlpack.c b/numpy/core/src/multiarray/dlpack.c index e4886cf4b..d5b1af101 100644 --- a/numpy/core/src/multiarray/dlpack.c +++ b/numpy/core/src/multiarray/dlpack.c @@ -269,7 +269,7 @@ array_dlpack_device(PyArrayObject *self, PyObject *NPY_UNUSED(args)) } NPY_NO_EXPORT PyObject * -_from_dlpack(PyObject *NPY_UNUSED(self), PyObject *obj) { +from_dlpack(PyObject *NPY_UNUSED(self), PyObject *obj) { PyObject *capsule = PyObject_CallMethod((PyObject *)obj->ob_type, "__dlpack__", "O", obj); if (capsule == NULL) { diff --git a/numpy/core/src/multiarray/item_selection.c b/numpy/core/src/multiarray/item_selection.c index 086b674c8..9fad153a3 100644 --- a/numpy/core/src/multiarray/item_selection.c +++ b/numpy/core/src/multiarray/item_selection.c @@ -2641,13 +2641,33 @@ PyArray_Nonzero(PyArrayObject *self) *multi_index++ = j++; } } + /* + * Fallback to a branchless strategy to avoid branch misprediction + * stalls that are very expensive on most modern processors. + */ else { - npy_intp j; - for (j = 0; j < count; ++j) { - if (*data != 0) { - *multi_index++ = j; - } + npy_intp *multi_index_end = multi_index + nonzero_count; + npy_intp j = 0; + + /* Manually unroll for GCC and maybe other compilers */ + while (multi_index + 4 < multi_index_end) { + *multi_index = j; + multi_index += data[0] != 0; + *multi_index = j + 1; + multi_index += data[stride] != 0; + *multi_index = j + 2; + multi_index += data[stride * 2] != 0; + *multi_index = j + 3; + multi_index += data[stride * 3] != 0; + data += stride * 4; + j += 4; + } + + while (multi_index < multi_index_end) { + *multi_index = j; + multi_index += *data != 0; data += stride; + ++j; } } } diff --git a/numpy/core/src/multiarray/multiarraymodule.c b/numpy/core/src/multiarray/multiarraymodule.c index 12923a6c6..f206ce2c6 100644 --- a/numpy/core/src/multiarray/multiarraymodule.c +++ b/numpy/core/src/multiarray/multiarraymodule.c @@ -4485,7 +4485,7 @@ static struct PyMethodDef array_module_methods[] = { {"_reload_guard", (PyCFunction)_reload_guard, METH_NOARGS, "Give a warning on reload and big warning in sub-interpreters."}, - {"_from_dlpack", (PyCFunction)_from_dlpack, + {"from_dlpack", (PyCFunction)from_dlpack, METH_O, NULL}, {NULL, NULL, 0, NULL} /* sentinel */ }; diff --git a/numpy/core/src/multiarray/scalarapi.c b/numpy/core/src/multiarray/scalarapi.c index 48d86caf1..85ef6d3eb 100644 --- a/numpy/core/src/multiarray/scalarapi.c +++ b/numpy/core/src/multiarray/scalarapi.c @@ -88,83 +88,12 @@ scalar_value(PyObject *scalar, PyArray_Descr *descr) } /* - * Must be a user-defined type --- check to see which - * scalar it inherits from. + * Must be a user defined type with an associated (registered) dtype. + * Thus, it cannot be flexible (user dtypes cannot be), so we can (and + * pretty much have no choice but to) assume the below logic always works. + * I.e. this assumes that the logic would also works for most of our types. */ -#define _CHK(cls) PyObject_IsInstance(scalar, \ - (PyObject *)&Py##cls##ArrType_Type) -#define _IFCASE(cls) if (_CHK(cls)) return &PyArrayScalar_VAL(scalar, cls) - - if (_CHK(Number)) { - if (_CHK(Integer)) { - if (_CHK(SignedInteger)) { - _IFCASE(Byte); - _IFCASE(Short); - _IFCASE(Int); - _IFCASE(Long); - _IFCASE(LongLong); - _IFCASE(Timedelta); - } - else { - /* Unsigned Integer */ - _IFCASE(UByte); - _IFCASE(UShort); - _IFCASE(UInt); - _IFCASE(ULong); - _IFCASE(ULongLong); - } - } - else { - /* Inexact */ - if (_CHK(Floating)) { - _IFCASE(Half); - _IFCASE(Float); - _IFCASE(Double); - _IFCASE(LongDouble); - } - else { - /*ComplexFloating */ - _IFCASE(CFloat); - _IFCASE(CDouble); - _IFCASE(CLongDouble); - } - } - } - else if (_CHK(Bool)) { - return &PyArrayScalar_VAL(scalar, Bool); - } - else if (_CHK(Datetime)) { - return &PyArrayScalar_VAL(scalar, Datetime); - } - else if (_CHK(Flexible)) { - if (_CHK(String)) { - return (void *)PyBytes_AS_STRING(scalar); - } - if (_CHK(Unicode)) { - /* Treat this the same as the NPY_UNICODE base class */ - - /* lazy initialization, to reduce the memory used by string scalars */ - if (PyArrayScalar_VAL(scalar, Unicode) == NULL) { - Py_UCS4 *raw_data = PyUnicode_AsUCS4Copy(scalar); - if (raw_data == NULL) { - return NULL; - } - PyArrayScalar_VAL(scalar, Unicode) = raw_data; - return (void *)raw_data; - } - return PyArrayScalar_VAL(scalar, Unicode); - } - if (_CHK(Void)) { - /* Note: no & needed here, so can't use _IFCASE */ - return PyArrayScalar_VAL(scalar, Void); - } - } - else { - _IFCASE(Object); - } - - /* * Use the alignment flag to figure out where the data begins * after a PyObject_HEAD @@ -177,8 +106,6 @@ scalar_value(PyObject *scalar, PyArray_Descr *descr) memloc = ((memloc + align - 1)/align)*align; } return (void *)memloc; -#undef _IFCASE -#undef _CHK } /*NUMPY_API diff --git a/numpy/core/src/npysort/x86-qsort.dispatch.c.src b/numpy/core/src/npysort/x86-qsort.dispatch.c.src deleted file mode 100644 index b93c737cb..000000000 --- a/numpy/core/src/npysort/x86-qsort.dispatch.c.src +++ /dev/null @@ -1,587 +0,0 @@ -/*@targets - * $maxopt $keep_baseline avx512_skx - */ -// policy $keep_baseline is used to avoid skip building avx512_skx -// when its part of baseline features (--cpu-baseline), since -// 'baseline' option isn't specified within targets. - -#include "x86-qsort.h" -#define NPY_NO_DEPRECATED_API NPY_API_VERSION - -#ifdef NPY_HAVE_AVX512_SKX -#include <immintrin.h> -#include "numpy/npy_math.h" -#include "npy_sort.h" -#include "simd/simd.h" - - -/* - * Quicksort using AVX-512 for int, uint32 and float. The ideas and code are - * based on these two research papers: - * (1) Fast and Robust Vectorized In-Place Sorting of Primitive Types - * https://drops.dagstuhl.de/opus/volltexte/2021/13775/ - * (2) A Novel Hybrid Quicksort Algorithm Vectorized using AVX-512 on Intel Skylake - * https://arxiv.org/pdf/1704.08579.pdf - * - * High level idea: Vectorize the quicksort partitioning using AVX-512 - * compressstore instructions. The algorithm to pick the pivot is to use median of - * 72 elements picked at random. If the array size is < 128, then use - * Bitonic sorting network. Good resource for bitonic sorting network: - * http://mitp-content-server.mit.edu:18180/books/content/sectbyfn?collid=books_pres_0&fn=Chapter%2027.pdf&id=8030 - * - * Refer to https://github.com/numpy/numpy/pull/20133#issuecomment-958110340 for - * potential problems when converting this code to universal intrinsics framework. - */ - -/* - * Constants used in sorting 16 elements in a ZMM registers. Based on Bitonic - * sorting network (see - * https://en.wikipedia.org/wiki/Bitonic_sorter#/media/File:BitonicSort.svg) - */ -#define NETWORK1 14,15,12,13,10,11,8,9,6,7,4,5,2,3,0,1 -#define NETWORK2 12,13,14,15,8,9,10,11,4,5,6,7,0,1,2,3 -#define NETWORK3 8,9,10,11,12,13,14,15,0,1,2,3,4,5,6,7 -#define NETWORK4 13,12,15,14,9,8,11,10,5,4,7,6,1,0,3,2 -#define NETWORK5 0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15 -#define NETWORK6 11,10,9,8,15,14,13,12,3,2,1,0,7,6,5,4 -#define NETWORK7 7,6,5,4,3,2,1,0,15,14,13,12,11,10,9,8 -#define ZMM_MAX_FLOAT _mm512_set1_ps(NPY_INFINITYF) -#define ZMM_MAX_UINT _mm512_set1_epi32(NPY_MAX_UINT32) -#define ZMM_MAX_INT _mm512_set1_epi32(NPY_MAX_INT32) -#define SHUFFLE_MASK(a,b,c,d) (a << 6) | (b << 4) | (c << 2) | d -#define SHUFFLE_ps(ZMM, MASK) _mm512_shuffle_ps(zmm, zmm, MASK) -#define SHUFFLE_epi32(ZMM, MASK) _mm512_shuffle_epi32(zmm, MASK) - -#define MAX(x, y) (((x) > (y)) ? (x) : (y)) -#define MIN(x, y) (((x) < (y)) ? (x) : (y)) - -/* - * Vectorized random number generator xoroshiro128+. Broken into 2 parts: - * (1) vnext generates 2 64-bit random integers - * (2) rnd_epu32 converts this to 4 32-bit random integers and bounds it to - * the length of the array - */ -#define VROTL(x, k) /* rotate each uint64_t value in vector */ \ - _mm256_or_si256(_mm256_slli_epi64((x),(k)),_mm256_srli_epi64((x),64-(k))) - -static NPY_INLINE -__m256i vnext(__m256i* s0, __m256i* s1) { - *s1 = _mm256_xor_si256(*s0, *s1); /* modify vectors s1 and s0 */ - *s0 = _mm256_xor_si256(_mm256_xor_si256(VROTL(*s0, 24), *s1), - _mm256_slli_epi64(*s1, 16)); - *s1 = VROTL(*s1, 37); - return _mm256_add_epi64(*s0, *s1); /* return random vector */ -} - -/* transform random numbers to the range between 0 and bound - 1 */ -static NPY_INLINE -__m256i rnd_epu32(__m256i rnd_vec, __m256i bound) { - __m256i even = _mm256_srli_epi64(_mm256_mul_epu32(rnd_vec, bound), 32); - __m256i odd = _mm256_mul_epu32(_mm256_srli_epi64(rnd_vec, 32), bound); - return _mm256_blend_epi32(odd, even, 0b01010101); -} - -/**begin repeat - * - * #TYPE = INT, UINT, FLOAT# - * #type = int, uint, float# - * #type_t = npy_int, npy_uint, npy_float# - * #zmm_t = __m512i, __m512i, __m512# - * #ymm_t = __m256i, __m256i, __m256# - * #vsuf1 = epi32, epu32, ps# - * #vsuf2 = epi32, epi32, ps# - * #vsuf3 = si512, si512, ps# - * #vsuf4 = s32, u32, f32# - * #CMP_GE_OP = _MM_CMPINT_NLT, _MM_CMPINT_NLT, _CMP_GE_OQ# - * #TYPE_MAX_VAL = NPY_MAX_INT32, NPY_MAX_UINT32, NPY_INFINITYF# - * #TYPE_MIN_VAL = NPY_MIN_INT32, 0, -NPY_INFINITYF# - */ - -/* - * COEX == Compare and Exchange two registers by swapping min and max values - */ -#define COEX_ZMM_@vsuf1@(a, b) { \ - @zmm_t@ temp = a; \ - a = _mm512_min_@vsuf1@(a,b); \ - b = _mm512_max_@vsuf1@(temp, b);} \ - -#define COEX_YMM_@vsuf1@(a, b){ \ - @ymm_t@ temp = a; \ - a = _mm256_min_@vsuf1@(a, b); \ - b = _mm256_max_@vsuf1@(temp, b);} \ - -static NPY_INLINE -@zmm_t@ cmp_merge_@vsuf1@(@zmm_t@ in1, @zmm_t@ in2, __mmask16 mask) -{ - @zmm_t@ min = _mm512_min_@vsuf1@(in2, in1); - @zmm_t@ max = _mm512_max_@vsuf1@(in2, in1); - return _mm512_mask_mov_@vsuf2@(min, mask, max); // 0 -> min, 1 -> max -} - -/* - * Assumes zmm is random and performs a full sorting network defined in - * https://en.wikipedia.org/wiki/Bitonic_sorter#/media/File:BitonicSort.svg - */ -static NPY_INLINE -@zmm_t@ sort_zmm_@vsuf1@(@zmm_t@ zmm) -{ - zmm = cmp_merge_@vsuf1@(zmm, SHUFFLE_@vsuf2@(zmm, SHUFFLE_MASK(2,3,0,1)), 0xAAAA); - zmm = cmp_merge_@vsuf1@(zmm, SHUFFLE_@vsuf2@(zmm, SHUFFLE_MASK(0,1,2,3)), 0xCCCC); - zmm = cmp_merge_@vsuf1@(zmm, SHUFFLE_@vsuf2@(zmm, SHUFFLE_MASK(2,3,0,1)), 0xAAAA); - zmm = cmp_merge_@vsuf1@(zmm, _mm512_permutexvar_@vsuf2@(_mm512_set_epi32(NETWORK3),zmm), 0xF0F0); - zmm = cmp_merge_@vsuf1@(zmm, SHUFFLE_@vsuf2@(zmm, SHUFFLE_MASK(1,0,3,2)), 0xCCCC); - zmm = cmp_merge_@vsuf1@(zmm, SHUFFLE_@vsuf2@(zmm, SHUFFLE_MASK(2,3,0,1)), 0xAAAA); - zmm = cmp_merge_@vsuf1@(zmm, _mm512_permutexvar_@vsuf2@(_mm512_set_epi32(NETWORK5),zmm), 0xFF00); - zmm = cmp_merge_@vsuf1@(zmm, _mm512_permutexvar_@vsuf2@(_mm512_set_epi32(NETWORK6),zmm), 0xF0F0); - zmm = cmp_merge_@vsuf1@(zmm, SHUFFLE_@vsuf2@(zmm, SHUFFLE_MASK(1,0,3,2)), 0xCCCC); - zmm = cmp_merge_@vsuf1@(zmm, SHUFFLE_@vsuf2@(zmm, SHUFFLE_MASK(2,3,0,1)), 0xAAAA); - return zmm; -} - -// Assumes zmm is bitonic and performs a recursive half cleaner -static NPY_INLINE -@zmm_t@ bitonic_merge_zmm_@vsuf1@(@zmm_t@ zmm) -{ - // 1) half_cleaner[16]: compare 1-9, 2-10, 3-11 etc .. - zmm = cmp_merge_@vsuf1@(zmm, _mm512_permutexvar_@vsuf2@(_mm512_set_epi32(NETWORK7),zmm), 0xFF00); - // 2) half_cleaner[8]: compare 1-5, 2-6, 3-7 etc .. - zmm = cmp_merge_@vsuf1@(zmm, _mm512_permutexvar_@vsuf2@(_mm512_set_epi32(NETWORK6),zmm), 0xF0F0); - // 3) half_cleaner[4] - zmm = cmp_merge_@vsuf1@(zmm, SHUFFLE_@vsuf2@(zmm, SHUFFLE_MASK(1,0,3,2)), 0xCCCC); - // 3) half_cleaner[1] - zmm = cmp_merge_@vsuf1@(zmm, SHUFFLE_@vsuf2@(zmm, SHUFFLE_MASK(2,3,0,1)), 0xAAAA); - return zmm; -} - -// Assumes zmm1 and zmm2 are sorted and performs a recursive half cleaner -static NPY_INLINE -void bitonic_merge_two_zmm_@vsuf1@(@zmm_t@* zmm1, @zmm_t@* zmm2) -{ - // 1) First step of a merging network: coex of zmm1 and zmm2 reversed - *zmm2 = _mm512_permutexvar_@vsuf2@(_mm512_set_epi32(NETWORK5), *zmm2); - @zmm_t@ zmm3 = _mm512_min_@vsuf1@(*zmm1, *zmm2); - @zmm_t@ zmm4 = _mm512_max_@vsuf1@(*zmm1, *zmm2); - // 2) Recursive half cleaner for each - *zmm1 = bitonic_merge_zmm_@vsuf1@(zmm3); - *zmm2 = bitonic_merge_zmm_@vsuf1@(zmm4); -} - -// Assumes [zmm0, zmm1] and [zmm2, zmm3] are sorted and performs a recursive half cleaner -static NPY_INLINE -void bitonic_merge_four_zmm_@vsuf1@(@zmm_t@* zmm) -{ - @zmm_t@ zmm2r = _mm512_permutexvar_@vsuf2@(_mm512_set_epi32(NETWORK5), zmm[2]); - @zmm_t@ zmm3r = _mm512_permutexvar_@vsuf2@(_mm512_set_epi32(NETWORK5), zmm[3]); - @zmm_t@ zmm_t1 = _mm512_min_@vsuf1@(zmm[0], zmm3r); - @zmm_t@ zmm_t2 = _mm512_min_@vsuf1@(zmm[1], zmm2r); - @zmm_t@ zmm_t3 = _mm512_permutexvar_@vsuf2@(_mm512_set_epi32(NETWORK5), _mm512_max_@vsuf1@(zmm[1], zmm2r)); - @zmm_t@ zmm_t4 = _mm512_permutexvar_@vsuf2@(_mm512_set_epi32(NETWORK5), _mm512_max_@vsuf1@(zmm[0], zmm3r)); - @zmm_t@ zmm0 = _mm512_min_@vsuf1@(zmm_t1, zmm_t2); - @zmm_t@ zmm1 = _mm512_max_@vsuf1@(zmm_t1, zmm_t2); - @zmm_t@ zmm2 = _mm512_min_@vsuf1@(zmm_t3, zmm_t4); - @zmm_t@ zmm3 = _mm512_max_@vsuf1@(zmm_t3, zmm_t4); - zmm[0] = bitonic_merge_zmm_@vsuf1@(zmm0); - zmm[1] = bitonic_merge_zmm_@vsuf1@(zmm1); - zmm[2] = bitonic_merge_zmm_@vsuf1@(zmm2); - zmm[3] = bitonic_merge_zmm_@vsuf1@(zmm3); -} - -static NPY_INLINE -void bitonic_merge_eight_zmm_@vsuf1@(@zmm_t@* zmm) -{ - @zmm_t@ zmm4r = _mm512_permutexvar_@vsuf2@(_mm512_set_epi32(NETWORK5), zmm[4]); - @zmm_t@ zmm5r = _mm512_permutexvar_@vsuf2@(_mm512_set_epi32(NETWORK5), zmm[5]); - @zmm_t@ zmm6r = _mm512_permutexvar_@vsuf2@(_mm512_set_epi32(NETWORK5), zmm[6]); - @zmm_t@ zmm7r = _mm512_permutexvar_@vsuf2@(_mm512_set_epi32(NETWORK5), zmm[7]); - @zmm_t@ zmm_t1 = _mm512_min_@vsuf1@(zmm[0], zmm7r); - @zmm_t@ zmm_t2 = _mm512_min_@vsuf1@(zmm[1], zmm6r); - @zmm_t@ zmm_t3 = _mm512_min_@vsuf1@(zmm[2], zmm5r); - @zmm_t@ zmm_t4 = _mm512_min_@vsuf1@(zmm[3], zmm4r); - @zmm_t@ zmm_t5 = _mm512_permutexvar_@vsuf2@(_mm512_set_epi32(NETWORK5), _mm512_max_@vsuf1@(zmm[3], zmm4r)); - @zmm_t@ zmm_t6 = _mm512_permutexvar_@vsuf2@(_mm512_set_epi32(NETWORK5), _mm512_max_@vsuf1@(zmm[2], zmm5r)); - @zmm_t@ zmm_t7 = _mm512_permutexvar_@vsuf2@(_mm512_set_epi32(NETWORK5), _mm512_max_@vsuf1@(zmm[1], zmm6r)); - @zmm_t@ zmm_t8 = _mm512_permutexvar_@vsuf2@(_mm512_set_epi32(NETWORK5), _mm512_max_@vsuf1@(zmm[0], zmm7r)); - COEX_ZMM_@vsuf1@(zmm_t1, zmm_t3); - COEX_ZMM_@vsuf1@(zmm_t2, zmm_t4); - COEX_ZMM_@vsuf1@(zmm_t5, zmm_t7); - COEX_ZMM_@vsuf1@(zmm_t6, zmm_t8); - COEX_ZMM_@vsuf1@(zmm_t1, zmm_t2); - COEX_ZMM_@vsuf1@(zmm_t3, zmm_t4); - COEX_ZMM_@vsuf1@(zmm_t5, zmm_t6); - COEX_ZMM_@vsuf1@(zmm_t7, zmm_t8); - zmm[0] = bitonic_merge_zmm_@vsuf1@(zmm_t1); - zmm[1] = bitonic_merge_zmm_@vsuf1@(zmm_t2); - zmm[2] = bitonic_merge_zmm_@vsuf1@(zmm_t3); - zmm[3] = bitonic_merge_zmm_@vsuf1@(zmm_t4); - zmm[4] = bitonic_merge_zmm_@vsuf1@(zmm_t5); - zmm[5] = bitonic_merge_zmm_@vsuf1@(zmm_t6); - zmm[6] = bitonic_merge_zmm_@vsuf1@(zmm_t7); - zmm[7] = bitonic_merge_zmm_@vsuf1@(zmm_t8); -} - -static NPY_INLINE -void sort_16_@vsuf1@(@type_t@* arr, npy_int N) -{ - __mmask16 load_mask = (0x0001 << N) - 0x0001; - @zmm_t@ zmm = _mm512_mask_loadu_@vsuf2@(ZMM_MAX_@TYPE@, load_mask, arr); - _mm512_mask_storeu_@vsuf2@(arr, load_mask, sort_zmm_@vsuf1@(zmm)); -} - -static NPY_INLINE -void sort_32_@vsuf1@(@type_t@* arr, npy_int N) -{ - if (N <= 16) { - sort_16_@vsuf1@(arr, N); - return; - } - @zmm_t@ zmm1 = _mm512_loadu_@vsuf3@(arr); - __mmask16 load_mask = (0x0001 << (N-16)) - 0x0001; - @zmm_t@ zmm2 = _mm512_mask_loadu_@vsuf2@(ZMM_MAX_@TYPE@, load_mask, arr + 16); - zmm1 = sort_zmm_@vsuf1@(zmm1); - zmm2 = sort_zmm_@vsuf1@(zmm2); - bitonic_merge_two_zmm_@vsuf1@(&zmm1, &zmm2); - _mm512_storeu_@vsuf3@(arr, zmm1); - _mm512_mask_storeu_@vsuf2@(arr + 16, load_mask, zmm2); -} - -static NPY_INLINE -void sort_64_@vsuf1@(@type_t@* arr, npy_int N) -{ - if (N <= 32) { - sort_32_@vsuf1@(arr, N); - return; - } - @zmm_t@ zmm[4]; - zmm[0] = _mm512_loadu_@vsuf3@(arr); - zmm[1] = _mm512_loadu_@vsuf3@(arr + 16); - __mmask16 load_mask1 = 0xFFFF, load_mask2 = 0xFFFF; - if (N < 48) { - load_mask1 = (0x0001 << (N-32)) - 0x0001; - load_mask2 = 0x0000; - } - else if (N < 64) { - load_mask2 = (0x0001 << (N-48)) - 0x0001; - } - zmm[2] = _mm512_mask_loadu_@vsuf2@(ZMM_MAX_@TYPE@, load_mask1, arr + 32); - zmm[3] = _mm512_mask_loadu_@vsuf2@(ZMM_MAX_@TYPE@, load_mask2, arr + 48); - zmm[0] = sort_zmm_@vsuf1@(zmm[0]); - zmm[1] = sort_zmm_@vsuf1@(zmm[1]); - zmm[2] = sort_zmm_@vsuf1@(zmm[2]); - zmm[3] = sort_zmm_@vsuf1@(zmm[3]); - bitonic_merge_two_zmm_@vsuf1@(&zmm[0], &zmm[1]); - bitonic_merge_two_zmm_@vsuf1@(&zmm[2], &zmm[3]); - bitonic_merge_four_zmm_@vsuf1@(zmm); - _mm512_storeu_@vsuf3@(arr, zmm[0]); - _mm512_storeu_@vsuf3@(arr + 16, zmm[1]); - _mm512_mask_storeu_@vsuf2@(arr + 32, load_mask1, zmm[2]); - _mm512_mask_storeu_@vsuf2@(arr + 48, load_mask2, zmm[3]); -} - -static NPY_INLINE -void sort_128_@vsuf1@(@type_t@* arr, npy_int N) -{ - if (N <= 64) { - sort_64_@vsuf1@(arr, N); - return; - } - @zmm_t@ zmm[8]; - zmm[0] = _mm512_loadu_@vsuf3@(arr); - zmm[1] = _mm512_loadu_@vsuf3@(arr + 16); - zmm[2] = _mm512_loadu_@vsuf3@(arr + 32); - zmm[3] = _mm512_loadu_@vsuf3@(arr + 48); - zmm[0] = sort_zmm_@vsuf1@(zmm[0]); - zmm[1] = sort_zmm_@vsuf1@(zmm[1]); - zmm[2] = sort_zmm_@vsuf1@(zmm[2]); - zmm[3] = sort_zmm_@vsuf1@(zmm[3]); - __mmask16 load_mask1 = 0xFFFF, load_mask2 = 0xFFFF; - __mmask16 load_mask3 = 0xFFFF, load_mask4 = 0xFFFF; - if (N < 80) { - load_mask1 = (0x0001 << (N-64)) - 0x0001; - load_mask2 = 0x0000; - load_mask3 = 0x0000; - load_mask4 = 0x0000; - } - else if (N < 96) { - load_mask2 = (0x0001 << (N-80)) - 0x0001; - load_mask3 = 0x0000; - load_mask4 = 0x0000; - } - else if (N < 112) { - load_mask3 = (0x0001 << (N-96)) - 0x0001; - load_mask4 = 0x0000; - } - else { - load_mask4 = (0x0001 << (N-112)) - 0x0001; - } - zmm[4] = _mm512_mask_loadu_@vsuf2@(ZMM_MAX_@TYPE@, load_mask1, arr + 64); - zmm[5] = _mm512_mask_loadu_@vsuf2@(ZMM_MAX_@TYPE@, load_mask2, arr + 80); - zmm[6] = _mm512_mask_loadu_@vsuf2@(ZMM_MAX_@TYPE@, load_mask3, arr + 96); - zmm[7] = _mm512_mask_loadu_@vsuf2@(ZMM_MAX_@TYPE@, load_mask4, arr + 112); - zmm[4] = sort_zmm_@vsuf1@(zmm[4]); - zmm[5] = sort_zmm_@vsuf1@(zmm[5]); - zmm[6] = sort_zmm_@vsuf1@(zmm[6]); - zmm[7] = sort_zmm_@vsuf1@(zmm[7]); - bitonic_merge_two_zmm_@vsuf1@(&zmm[0], &zmm[1]); - bitonic_merge_two_zmm_@vsuf1@(&zmm[2], &zmm[3]); - bitonic_merge_two_zmm_@vsuf1@(&zmm[4], &zmm[5]); - bitonic_merge_two_zmm_@vsuf1@(&zmm[6], &zmm[7]); - bitonic_merge_four_zmm_@vsuf1@(zmm); - bitonic_merge_four_zmm_@vsuf1@(zmm + 4); - bitonic_merge_eight_zmm_@vsuf1@(zmm); - _mm512_storeu_@vsuf3@(arr, zmm[0]); - _mm512_storeu_@vsuf3@(arr + 16, zmm[1]); - _mm512_storeu_@vsuf3@(arr + 32, zmm[2]); - _mm512_storeu_@vsuf3@(arr + 48, zmm[3]); - _mm512_mask_storeu_@vsuf2@(arr + 64, load_mask1, zmm[4]); - _mm512_mask_storeu_@vsuf2@(arr + 80, load_mask2, zmm[5]); - _mm512_mask_storeu_@vsuf2@(arr + 96, load_mask3, zmm[6]); - _mm512_mask_storeu_@vsuf2@(arr + 112, load_mask4, zmm[7]); -} - - -static NPY_INLINE -void swap_@TYPE@(@type_t@ *arr, npy_intp ii, npy_intp jj) { - @type_t@ temp = arr[ii]; - arr[ii] = arr[jj]; - arr[jj] = temp; -} - -// Median of 3 stratergy -//static NPY_INLINE -//npy_intp get_pivot_index(@type_t@ *arr, const npy_intp left, const npy_intp right) { -// return (rand() % (right + 1 - left)) + left; -// //npy_intp middle = ((right-left)/2) + left; -// //@type_t@ a = arr[left], b = arr[middle], c = arr[right]; -// //if ((b >= a && b <= c) || (b <= a && b >= c)) -// // return middle; -// //if ((a >= b && a <= c) || (a <= b && a >= c)) -// // return left; -// //else -// // return right; -//} - -/* - * Picking the pivot: Median of 72 array elements chosen at random. - */ - -static NPY_INLINE -@type_t@ get_pivot_@vsuf1@(@type_t@ *arr, const npy_intp left, const npy_intp right) { - /* seeds for vectorized random number generator */ - __m256i s0 = _mm256_setr_epi64x(8265987198341093849, 3762817312854612374, - 1324281658759788278, 6214952190349879213); - __m256i s1 = _mm256_setr_epi64x(2874178529384792648, 1257248936691237653, - 7874578921548791257, 1998265912745817298); - s0 = _mm256_add_epi64(s0, _mm256_set1_epi64x(left)); - s1 = _mm256_sub_epi64(s1, _mm256_set1_epi64x(right)); - - npy_intp arrsize = right - left + 1; - __m256i bound = _mm256_set1_epi32(arrsize > INT32_MAX ? INT32_MAX : arrsize); - __m512i left_vec = _mm512_set1_epi64(left); - __m512i right_vec = _mm512_set1_epi64(right); - @ymm_t@ v[9]; - /* fill 9 vectors with random numbers */ - for (npy_int i = 0; i < 9; ++i) { - __m256i rand_64 = vnext(&s0, &s1); /* vector with 4 random uint64_t */ - __m512i rand_32 = _mm512_cvtepi32_epi64(rnd_epu32(rand_64, bound)); /* random numbers between 0 and bound - 1 */ - __m512i indices; - if (i < 5) - indices = _mm512_add_epi64(left_vec, rand_32); /* indices for arr */ - else - indices = _mm512_sub_epi64(right_vec, rand_32); /* indices for arr */ - - v[i] = _mm512_i64gather_@vsuf2@(indices, arr, sizeof(@type_t@)); - } - - /* median network for 9 elements */ - COEX_YMM_@vsuf1@(v[0], v[1]); COEX_YMM_@vsuf1@(v[2], v[3]); - COEX_YMM_@vsuf1@(v[4], v[5]); COEX_YMM_@vsuf1@(v[6], v[7]); - COEX_YMM_@vsuf1@(v[0], v[2]); COEX_YMM_@vsuf1@(v[1], v[3]); - COEX_YMM_@vsuf1@(v[4], v[6]); COEX_YMM_@vsuf1@(v[5], v[7]); - COEX_YMM_@vsuf1@(v[0], v[4]); COEX_YMM_@vsuf1@(v[1], v[2]); - COEX_YMM_@vsuf1@(v[5], v[6]); COEX_YMM_@vsuf1@(v[3], v[7]); - COEX_YMM_@vsuf1@(v[1], v[5]); COEX_YMM_@vsuf1@(v[2], v[6]); - COEX_YMM_@vsuf1@(v[3], v[5]); COEX_YMM_@vsuf1@(v[2], v[4]); - COEX_YMM_@vsuf1@(v[3], v[4]); - COEX_YMM_@vsuf1@(v[3], v[8]); - COEX_YMM_@vsuf1@(v[4], v[8]); - - // technically v[4] needs to be sorted before we pick the correct median, - // picking the 4th element works just as well for performance - @type_t@* temp = (@type_t@*) &v[4]; - - return temp[4]; -} - -/* - * Parition one ZMM register based on the pivot and returns the index of the - * last element that is less than equal to the pivot. - */ -static NPY_INLINE -npy_int partition_vec_@vsuf1@(@type_t@* arr, npy_intp left, npy_intp right, - const @zmm_t@ curr_vec, const @zmm_t@ pivot_vec, - @zmm_t@* smallest_vec, @zmm_t@* biggest_vec) -{ - /* which elements are larger than the pivot */ - __mmask16 gt_mask = _mm512_cmp_@vsuf1@_mask(curr_vec, pivot_vec, @CMP_GE_OP@); - npy_int amount_gt_pivot = _mm_popcnt_u32((npy_int)gt_mask); - _mm512_mask_compressstoreu_@vsuf2@(arr + left, _knot_mask16(gt_mask), curr_vec); - _mm512_mask_compressstoreu_@vsuf2@(arr + right - amount_gt_pivot, gt_mask, curr_vec); - *smallest_vec = _mm512_min_@vsuf1@(curr_vec, *smallest_vec); - *biggest_vec = _mm512_max_@vsuf1@(curr_vec, *biggest_vec); - return amount_gt_pivot; -} - -/* - * Parition an array based on the pivot and returns the index of the - * last element that is less than equal to the pivot. - */ -static NPY_INLINE -npy_intp partition_avx512_@vsuf1@(@type_t@* arr, npy_intp left, npy_intp right, - @type_t@ pivot, @type_t@* smallest, @type_t@* biggest) -{ - /* make array length divisible by 16 , shortening the array */ - for (npy_int i = (right - left) % 16; i > 0; --i) { - *smallest = MIN(*smallest, arr[left]); - *biggest = MAX(*biggest, arr[left]); - if (arr[left] > pivot) { - swap_@TYPE@(arr, left, --right); - } - else { - ++left; - } - } - - if(left == right) - return left; /* less than 16 elements in the array */ - - @zmm_t@ pivot_vec = _mm512_set1_@vsuf2@(pivot); - @zmm_t@ min_vec = _mm512_set1_@vsuf2@(*smallest); - @zmm_t@ max_vec = _mm512_set1_@vsuf2@(*biggest); - - if(right - left == 16) { - @zmm_t@ vec = _mm512_loadu_@vsuf3@(arr + left); - npy_int amount_gt_pivot = partition_vec_@vsuf1@(arr, left, left + 16, vec, pivot_vec, &min_vec, &max_vec); - *smallest = npyv_reducemin_@vsuf4@(min_vec); - *biggest = npyv_reducemax_@vsuf4@(max_vec); - return left + (16 - amount_gt_pivot); - } - - // first and last 16 values are partitioned at the end - @zmm_t@ vec_left = _mm512_loadu_@vsuf3@(arr + left); - @zmm_t@ vec_right = _mm512_loadu_@vsuf3@(arr + (right - 16)); - // store points of the vectors - npy_intp r_store = right - 16; - npy_intp l_store = left; - // indices for loading the elements - left += 16; - right -= 16; - while(right - left != 0) { - @zmm_t@ curr_vec; - /* - * if fewer elements are stored on the right side of the array, - * then next elements are loaded from the right side, - * otherwise from the left side - */ - if((r_store + 16) - right < left - l_store) { - right -= 16; - curr_vec = _mm512_loadu_@vsuf3@(arr + right); - } - else { - curr_vec = _mm512_loadu_@vsuf3@(arr + left); - left += 16; - } - // partition the current vector and save it on both sides of the array - npy_int amount_gt_pivot = partition_vec_@vsuf1@(arr, l_store, r_store + 16, curr_vec, pivot_vec, &min_vec, &max_vec);; - r_store -= amount_gt_pivot; l_store += (16 - amount_gt_pivot); - } - - /* partition and save vec_left and vec_right */ - npy_int amount_gt_pivot = partition_vec_@vsuf1@(arr, l_store, r_store + 16, vec_left, pivot_vec, &min_vec, &max_vec); - l_store += (16 - amount_gt_pivot); - amount_gt_pivot = partition_vec_@vsuf1@(arr, l_store, l_store + 16, vec_right, pivot_vec, &min_vec, &max_vec); - l_store += (16 - amount_gt_pivot); - *smallest = npyv_reducemin_@vsuf4@(min_vec); - *biggest = npyv_reducemax_@vsuf4@(max_vec); - return l_store; -} - -static NPY_INLINE -void qsort_@type@(@type_t@* arr, npy_intp left, npy_intp right, npy_int max_iters) -{ - /* - * Resort to heapsort if quicksort isnt making any progress - */ - if (max_iters <= 0) { - heapsort_@type@((void*)(arr + left), right + 1 - left, NULL); - return; - } - /* - * Base case: use bitonic networks to sort arrays <= 128 - */ - if (right + 1 - left <= 128) { - sort_128_@vsuf1@(arr + left, right + 1 -left); - return; - } - - @type_t@ pivot = get_pivot_@vsuf1@(arr, left, right); - @type_t@ smallest = @TYPE_MAX_VAL@; - @type_t@ biggest = @TYPE_MIN_VAL@; - npy_intp pivot_index = partition_avx512_@vsuf1@(arr, left, right+1, pivot, &smallest, &biggest); - if (pivot != smallest) - qsort_@type@(arr, left, pivot_index - 1, max_iters - 1); - if (pivot != biggest) - qsort_@type@(arr, pivot_index, right, max_iters - 1); -} -/**end repeat**/ - -static NPY_INLINE -npy_intp replace_nan_with_inf(npy_float* arr, npy_intp arrsize) -{ - npy_intp nan_count = 0; - __mmask16 loadmask = 0xFFFF; - while (arrsize > 0) { - if (arrsize < 16) { - loadmask = (0x0001 << arrsize) - 0x0001; - } - __m512 in_zmm = _mm512_maskz_loadu_ps(loadmask, arr); - __mmask16 nanmask = _mm512_cmp_ps_mask(in_zmm, in_zmm, _CMP_NEQ_UQ); - nan_count += _mm_popcnt_u32((npy_int) nanmask); - _mm512_mask_storeu_ps(arr, nanmask, ZMM_MAX_FLOAT); - arr += 16; - arrsize -= 16; - } - return nan_count; -} - -static NPY_INLINE -void replace_inf_with_nan(npy_float* arr, npy_intp arrsize, npy_intp nan_count) -{ - for (npy_intp ii = arrsize-1; nan_count > 0; --ii) { - arr[ii] = NPY_NANF; - nan_count -= 1; - } -} - -/**begin repeat - * - * #type = int, uint, float# - * #type_t = npy_int, npy_uint, npy_float# - * #FIXNAN = 0, 0, 1# - */ - -NPY_NO_EXPORT void NPY_CPU_DISPATCH_CURFX(x86_quicksort_@type@) -(void* arr, npy_intp arrsize) -{ - if (arrsize > 1) { -#if @FIXNAN@ - npy_intp nan_count = replace_nan_with_inf((@type_t@*) arr, arrsize); -#endif - qsort_@type@((@type_t@*) arr, 0, arrsize-1, 2*log2(arrsize)); -#if @FIXNAN@ - replace_inf_with_nan((@type_t@*) arr, arrsize, nan_count); -#endif - } -} -/**end repeat**/ - -#endif // NPY_HAVE_AVX512_SKX diff --git a/numpy/core/src/npysort/x86-qsort.dispatch.cpp b/numpy/core/src/npysort/x86-qsort.dispatch.cpp new file mode 100644 index 000000000..4b01e3528 --- /dev/null +++ b/numpy/core/src/npysort/x86-qsort.dispatch.cpp @@ -0,0 +1,835 @@ +/*@targets + * $maxopt $keep_baseline avx512_skx + */ +// policy $keep_baseline is used to avoid skip building avx512_skx +// when its part of baseline features (--cpu-baseline), since +// 'baseline' option isn't specified within targets. + +#include "x86-qsort.h" +#define NPY_NO_DEPRECATED_API NPY_API_VERSION + +#ifdef NPY_HAVE_AVX512_SKX +#include "numpy/npy_math.h" + +#include "npy_sort.h" +#include "numpy_tag.h" + +#include "simd/simd.h" +#include <immintrin.h> + +template <typename Tag, typename type> +NPY_NO_EXPORT int +heapsort_(type *start, npy_intp n); + +/* + * Quicksort using AVX-512 for int, uint32 and float. The ideas and code are + * based on these two research papers: + * (1) Fast and Robust Vectorized In-Place Sorting of Primitive Types + * https://drops.dagstuhl.de/opus/volltexte/2021/13775/ + * (2) A Novel Hybrid Quicksort Algorithm Vectorized using AVX-512 on Intel + * Skylake https://arxiv.org/pdf/1704.08579.pdf + * + * High level idea: Vectorize the quicksort partitioning using AVX-512 + * compressstore instructions. The algorithm to pick the pivot is to use median + * of 72 elements picked at random. If the array size is < 128, then use + * Bitonic sorting network. Good resource for bitonic sorting network: + * http://mitp-content-server.mit.edu:18180/books/content/sectbyfn?collid=books_pres_0&fn=Chapter%2027.pdf&id=8030 + * + * Refer to https://github.com/numpy/numpy/pull/20133#issuecomment-958110340 + * for potential problems when converting this code to universal intrinsics + * framework. + */ + +/* + * Constants used in sorting 16 elements in a ZMM registers. Based on Bitonic + * sorting network (see + * https://en.wikipedia.org/wiki/Bitonic_sorter#/media/File:BitonicSort.svg) + */ +#define NETWORK1 14, 15, 12, 13, 10, 11, 8, 9, 6, 7, 4, 5, 2, 3, 0, 1 +#define NETWORK2 12, 13, 14, 15, 8, 9, 10, 11, 4, 5, 6, 7, 0, 1, 2, 3 +#define NETWORK3 8, 9, 10, 11, 12, 13, 14, 15, 0, 1, 2, 3, 4, 5, 6, 7 +#define NETWORK4 13, 12, 15, 14, 9, 8, 11, 10, 5, 4, 7, 6, 1, 0, 3, 2 +#define NETWORK5 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15 +#define NETWORK6 11, 10, 9, 8, 15, 14, 13, 12, 3, 2, 1, 0, 7, 6, 5, 4 +#define NETWORK7 7, 6, 5, 4, 3, 2, 1, 0, 15, 14, 13, 12, 11, 10, 9, 8 +#define ZMM_MAX_FLOAT _mm512_set1_ps(NPY_INFINITYF) +#define ZMM_MAX_UINT _mm512_set1_epi32(NPY_MAX_UINT32) +#define ZMM_MAX_INT _mm512_set1_epi32(NPY_MAX_INT32) +#define SHUFFLE_MASK(a, b, c, d) (a << 6) | (b << 4) | (c << 2) | d +#define SHUFFLE_ps(ZMM, MASK) _mm512_shuffle_ps(zmm, zmm, MASK) +#define SHUFFLE_epi32(ZMM, MASK) _mm512_shuffle_epi32(zmm, MASK) + +#define MAX(x, y) (((x) > (y)) ? (x) : (y)) +#define MIN(x, y) (((x) < (y)) ? (x) : (y)) + +/* + * Vectorized random number generator xoroshiro128+. Broken into 2 parts: + * (1) vnext generates 2 64-bit random integers + * (2) rnd_epu32 converts this to 4 32-bit random integers and bounds it to + * the length of the array + */ +#define VROTL(x, k) /* rotate each uint64_t value in vector */ \ + _mm256_or_si256(_mm256_slli_epi64((x), (k)), \ + _mm256_srli_epi64((x), 64 - (k))) + +static NPY_INLINE __m256i +vnext(__m256i *s0, __m256i *s1) +{ + *s1 = _mm256_xor_si256(*s0, *s1); /* modify vectors s1 and s0 */ + *s0 = _mm256_xor_si256(_mm256_xor_si256(VROTL(*s0, 24), *s1), + _mm256_slli_epi64(*s1, 16)); + *s1 = VROTL(*s1, 37); + return _mm256_add_epi64(*s0, *s1); /* return random vector */ +} + +/* transform random numbers to the range between 0 and bound - 1 */ +static NPY_INLINE __m256i +rnd_epu32(__m256i rnd_vec, __m256i bound) +{ + __m256i even = _mm256_srli_epi64(_mm256_mul_epu32(rnd_vec, bound), 32); + __m256i odd = _mm256_mul_epu32(_mm256_srli_epi64(rnd_vec, 32), bound); + return _mm256_blend_epi32(odd, even, 0b01010101); +} + +template <typename type> +struct vector; + +template <> +struct vector<npy_int> { + using tag = npy::int_tag; + using type_t = npy_int; + using zmm_t = __m512i; + using ymm_t = __m256i; + + static type_t type_max() { return NPY_MAX_INT32; } + static type_t type_min() { return NPY_MIN_INT32; } + static zmm_t zmm_max() { return _mm512_set1_epi32(type_max()); } + + static __mmask16 ge(zmm_t x, zmm_t y) + { + return _mm512_cmp_epi32_mask(x, y, _MM_CMPINT_NLT); + } + template <int scale> + static ymm_t i64gather(__m512i index, void const *base) + { + return _mm512_i64gather_epi32(index, base, scale); + } + static zmm_t loadu(void const *mem) { return _mm512_loadu_si512(mem); } + static zmm_t max(zmm_t x, zmm_t y) { return _mm512_max_epi32(x, y); } + static void mask_compressstoreu(void *mem, __mmask16 mask, zmm_t x) + { + return _mm512_mask_compressstoreu_epi32(mem, mask, x); + } + static zmm_t mask_loadu(zmm_t x, __mmask16 mask, void const *mem) + { + return _mm512_mask_loadu_epi32(x, mask, mem); + } + static zmm_t mask_mov(zmm_t x, __mmask16 mask, zmm_t y) + { + return _mm512_mask_mov_epi32(x, mask, y); + } + static void mask_storeu(void *mem, __mmask16 mask, zmm_t x) + { + return _mm512_mask_storeu_epi32(mem, mask, x); + } + static zmm_t min(zmm_t x, zmm_t y) { return _mm512_min_epi32(x, y); } + static zmm_t permutexvar(__m512i idx, zmm_t zmm) + { + return _mm512_permutexvar_epi32(idx, zmm); + } + static type_t reducemax(zmm_t v) { return npyv_reducemax_s32(v); } + static type_t reducemin(zmm_t v) { return npyv_reducemin_s32(v); } + static zmm_t set1(type_t v) { return _mm512_set1_epi32(v); } + template<__mmask16 mask> + static zmm_t shuffle(zmm_t zmm) + { + return _mm512_shuffle_epi32(zmm, (_MM_PERM_ENUM)mask); + } + static void storeu(void *mem, zmm_t x) + { + return _mm512_storeu_si512(mem, x); + } + + static ymm_t max(ymm_t x, ymm_t y) { return _mm256_max_epi32(x, y); } + static ymm_t min(ymm_t x, ymm_t y) { return _mm256_min_epi32(x, y); } +}; +template <> +struct vector<npy_uint> { + using tag = npy::uint_tag; + using type_t = npy_uint; + using zmm_t = __m512i; + using ymm_t = __m256i; + + static type_t type_max() { return NPY_MAX_UINT32; } + static type_t type_min() { return 0; } + static zmm_t zmm_max() { return _mm512_set1_epi32(type_max()); } + + template<int scale> + static ymm_t i64gather(__m512i index, void const *base) + { + return _mm512_i64gather_epi32(index, base, scale); + } + static __mmask16 ge(zmm_t x, zmm_t y) + { + return _mm512_cmp_epu32_mask(x, y, _MM_CMPINT_NLT); + } + static zmm_t loadu(void const *mem) { return _mm512_loadu_si512(mem); } + static zmm_t max(zmm_t x, zmm_t y) { return _mm512_max_epu32(x, y); } + static void mask_compressstoreu(void *mem, __mmask16 mask, zmm_t x) + { + return _mm512_mask_compressstoreu_epi32(mem, mask, x); + } + static zmm_t mask_loadu(zmm_t x, __mmask16 mask, void const *mem) + { + return _mm512_mask_loadu_epi32(x, mask, mem); + } + static zmm_t mask_mov(zmm_t x, __mmask16 mask, zmm_t y) + { + return _mm512_mask_mov_epi32(x, mask, y); + } + static void mask_storeu(void *mem, __mmask16 mask, zmm_t x) + { + return _mm512_mask_storeu_epi32(mem, mask, x); + } + static zmm_t min(zmm_t x, zmm_t y) { return _mm512_min_epu32(x, y); } + static zmm_t permutexvar(__m512i idx, zmm_t zmm) + { + return _mm512_permutexvar_epi32(idx, zmm); + } + static type_t reducemax(zmm_t v) { return npyv_reducemax_u32(v); } + static type_t reducemin(zmm_t v) { return npyv_reducemin_u32(v); } + static zmm_t set1(type_t v) { return _mm512_set1_epi32(v); } + template<__mmask16 mask> + static zmm_t shuffle(zmm_t zmm) + { + return _mm512_shuffle_epi32(zmm, (_MM_PERM_ENUM)mask); + } + static void storeu(void *mem, zmm_t x) + { + return _mm512_storeu_si512(mem, x); + } + + static ymm_t max(ymm_t x, ymm_t y) { return _mm256_max_epu32(x, y); } + static ymm_t min(ymm_t x, ymm_t y) { return _mm256_min_epu32(x, y); } +}; +template <> +struct vector<npy_float> { + using tag = npy::float_tag; + using type_t = npy_float; + using zmm_t = __m512; + using ymm_t = __m256; + + static type_t type_max() { return NPY_INFINITYF; } + static type_t type_min() { return -NPY_INFINITYF; } + static zmm_t zmm_max() { return _mm512_set1_ps(type_max()); } + + static __mmask16 ge(zmm_t x, zmm_t y) + { + return _mm512_cmp_ps_mask(x, y, _CMP_GE_OQ); + } + template<int scale> + static ymm_t i64gather(__m512i index, void const *base) + { + return _mm512_i64gather_ps(index, base, scale); + } + static zmm_t loadu(void const *mem) { return _mm512_loadu_ps(mem); } + static zmm_t max(zmm_t x, zmm_t y) { return _mm512_max_ps(x, y); } + static void mask_compressstoreu(void *mem, __mmask16 mask, zmm_t x) + { + return _mm512_mask_compressstoreu_ps(mem, mask, x); + } + static zmm_t mask_loadu(zmm_t x, __mmask16 mask, void const *mem) + { + return _mm512_mask_loadu_ps(x, mask, mem); + } + static zmm_t mask_mov(zmm_t x, __mmask16 mask, zmm_t y) + { + return _mm512_mask_mov_ps(x, mask, y); + } + static void mask_storeu(void *mem, __mmask16 mask, zmm_t x) + { + return _mm512_mask_storeu_ps(mem, mask, x); + } + static zmm_t min(zmm_t x, zmm_t y) { return _mm512_min_ps(x, y); } + static zmm_t permutexvar(__m512i idx, zmm_t zmm) + { + return _mm512_permutexvar_ps(idx, zmm); + } + static type_t reducemax(zmm_t v) { return npyv_reducemax_f32(v); } + static type_t reducemin(zmm_t v) { return npyv_reducemin_f32(v); } + static zmm_t set1(type_t v) { return _mm512_set1_ps(v); } + template<__mmask16 mask> + static zmm_t shuffle(zmm_t zmm) + { + return _mm512_shuffle_ps(zmm, zmm, (_MM_PERM_ENUM)mask); + } + static void storeu(void *mem, zmm_t x) { return _mm512_storeu_ps(mem, x); } + + static ymm_t max(ymm_t x, ymm_t y) { return _mm256_max_ps(x, y); } + static ymm_t min(ymm_t x, ymm_t y) { return _mm256_min_ps(x, y); } +}; + +/* + * COEX == Compare and Exchange two registers by swapping min and max values + */ +template <typename vtype, typename mm_t> +void +COEX(mm_t &a, mm_t &b) +{ + mm_t temp = a; + a = vtype::min(a, b); + b = vtype::max(temp, b); +} + +template <typename vtype, typename zmm_t = typename vtype::zmm_t> +static NPY_INLINE zmm_t +cmp_merge(zmm_t in1, zmm_t in2, __mmask16 mask) +{ + zmm_t min = vtype::min(in2, in1); + zmm_t max = vtype::max(in2, in1); + return vtype::mask_mov(min, mask, max); // 0 -> min, 1 -> max +} + +/* + * Assumes zmm is random and performs a full sorting network defined in + * https://en.wikipedia.org/wiki/Bitonic_sorter#/media/File:BitonicSort.svg + */ +template <typename vtype, typename zmm_t = typename vtype::zmm_t> +static NPY_INLINE zmm_t +sort_zmm(zmm_t zmm) +{ + zmm = cmp_merge<vtype>(zmm, vtype::template shuffle<SHUFFLE_MASK(2, 3, 0, 1)>(zmm), + 0xAAAA); + zmm = cmp_merge<vtype>(zmm, vtype::template shuffle<SHUFFLE_MASK(0, 1, 2, 3)>(zmm), + 0xCCCC); + zmm = cmp_merge<vtype>(zmm, vtype::template shuffle<SHUFFLE_MASK(2, 3, 0, 1)>(zmm), + 0xAAAA); + zmm = cmp_merge<vtype>( + zmm, vtype::permutexvar(_mm512_set_epi32(NETWORK3), zmm), 0xF0F0); + zmm = cmp_merge<vtype>(zmm, vtype::template shuffle<SHUFFLE_MASK(1, 0, 3, 2)>(zmm), + 0xCCCC); + zmm = cmp_merge<vtype>(zmm, vtype::template shuffle<SHUFFLE_MASK(2, 3, 0, 1)>(zmm), + 0xAAAA); + zmm = cmp_merge<vtype>( + zmm, vtype::permutexvar(_mm512_set_epi32(NETWORK5), zmm), 0xFF00); + zmm = cmp_merge<vtype>( + zmm, vtype::permutexvar(_mm512_set_epi32(NETWORK6), zmm), 0xF0F0); + zmm = cmp_merge<vtype>(zmm, vtype::template shuffle<SHUFFLE_MASK(1, 0, 3, 2)>(zmm), + 0xCCCC); + zmm = cmp_merge<vtype>(zmm, vtype::template shuffle<SHUFFLE_MASK(2, 3, 0, 1)>(zmm), + 0xAAAA); + return zmm; +} + +// Assumes zmm is bitonic and performs a recursive half cleaner +template <typename vtype, typename zmm_t = typename vtype::zmm_t> +static NPY_INLINE zmm_t +bitonic_merge_zmm(zmm_t zmm) +{ + // 1) half_cleaner[16]: compare 1-9, 2-10, 3-11 etc .. + zmm = cmp_merge<vtype>( + zmm, vtype::permutexvar(_mm512_set_epi32(NETWORK7), zmm), 0xFF00); + // 2) half_cleaner[8]: compare 1-5, 2-6, 3-7 etc .. + zmm = cmp_merge<vtype>( + zmm, vtype::permutexvar(_mm512_set_epi32(NETWORK6), zmm), 0xF0F0); + // 3) half_cleaner[4] + zmm = cmp_merge<vtype>(zmm, vtype::template shuffle<SHUFFLE_MASK(1, 0, 3, 2)>(zmm), + 0xCCCC); + // 3) half_cleaner[1] + zmm = cmp_merge<vtype>(zmm, vtype::template shuffle<SHUFFLE_MASK(2, 3, 0, 1)>(zmm), + 0xAAAA); + return zmm; +} + +// Assumes zmm1 and zmm2 are sorted and performs a recursive half cleaner +template <typename vtype, typename zmm_t = typename vtype::zmm_t> +static NPY_INLINE void +bitonic_merge_two_zmm(zmm_t *zmm1, zmm_t *zmm2) +{ + // 1) First step of a merging network: coex of zmm1 and zmm2 reversed + *zmm2 = vtype::permutexvar(_mm512_set_epi32(NETWORK5), *zmm2); + zmm_t zmm3 = vtype::min(*zmm1, *zmm2); + zmm_t zmm4 = vtype::max(*zmm1, *zmm2); + // 2) Recursive half cleaner for each + *zmm1 = bitonic_merge_zmm<vtype>(zmm3); + *zmm2 = bitonic_merge_zmm<vtype>(zmm4); +} + +// Assumes [zmm0, zmm1] and [zmm2, zmm3] are sorted and performs a recursive +// half cleaner +template <typename vtype, typename zmm_t = typename vtype::zmm_t> +static NPY_INLINE void +bitonic_merge_four_zmm(zmm_t *zmm) +{ + zmm_t zmm2r = vtype::permutexvar(_mm512_set_epi32(NETWORK5), zmm[2]); + zmm_t zmm3r = vtype::permutexvar(_mm512_set_epi32(NETWORK5), zmm[3]); + zmm_t zmm_t1 = vtype::min(zmm[0], zmm3r); + zmm_t zmm_t2 = vtype::min(zmm[1], zmm2r); + zmm_t zmm_t3 = vtype::permutexvar(_mm512_set_epi32(NETWORK5), + vtype::max(zmm[1], zmm2r)); + zmm_t zmm_t4 = vtype::permutexvar(_mm512_set_epi32(NETWORK5), + vtype::max(zmm[0], zmm3r)); + zmm_t zmm0 = vtype::min(zmm_t1, zmm_t2); + zmm_t zmm1 = vtype::max(zmm_t1, zmm_t2); + zmm_t zmm2 = vtype::min(zmm_t3, zmm_t4); + zmm_t zmm3 = vtype::max(zmm_t3, zmm_t4); + zmm[0] = bitonic_merge_zmm<vtype>(zmm0); + zmm[1] = bitonic_merge_zmm<vtype>(zmm1); + zmm[2] = bitonic_merge_zmm<vtype>(zmm2); + zmm[3] = bitonic_merge_zmm<vtype>(zmm3); +} + +template <typename vtype, typename zmm_t = typename vtype::zmm_t> +static NPY_INLINE void +bitonic_merge_eight_zmm(zmm_t *zmm) +{ + zmm_t zmm4r = vtype::permutexvar(_mm512_set_epi32(NETWORK5), zmm[4]); + zmm_t zmm5r = vtype::permutexvar(_mm512_set_epi32(NETWORK5), zmm[5]); + zmm_t zmm6r = vtype::permutexvar(_mm512_set_epi32(NETWORK5), zmm[6]); + zmm_t zmm7r = vtype::permutexvar(_mm512_set_epi32(NETWORK5), zmm[7]); + zmm_t zmm_t1 = vtype::min(zmm[0], zmm7r); + zmm_t zmm_t2 = vtype::min(zmm[1], zmm6r); + zmm_t zmm_t3 = vtype::min(zmm[2], zmm5r); + zmm_t zmm_t4 = vtype::min(zmm[3], zmm4r); + zmm_t zmm_t5 = vtype::permutexvar(_mm512_set_epi32(NETWORK5), + vtype::max(zmm[3], zmm4r)); + zmm_t zmm_t6 = vtype::permutexvar(_mm512_set_epi32(NETWORK5), + vtype::max(zmm[2], zmm5r)); + zmm_t zmm_t7 = vtype::permutexvar(_mm512_set_epi32(NETWORK5), + vtype::max(zmm[1], zmm6r)); + zmm_t zmm_t8 = vtype::permutexvar(_mm512_set_epi32(NETWORK5), + vtype::max(zmm[0], zmm7r)); + COEX<vtype>(zmm_t1, zmm_t3); + COEX<vtype>(zmm_t2, zmm_t4); + COEX<vtype>(zmm_t5, zmm_t7); + COEX<vtype>(zmm_t6, zmm_t8); + COEX<vtype>(zmm_t1, zmm_t2); + COEX<vtype>(zmm_t3, zmm_t4); + COEX<vtype>(zmm_t5, zmm_t6); + COEX<vtype>(zmm_t7, zmm_t8); + zmm[0] = bitonic_merge_zmm<vtype>(zmm_t1); + zmm[1] = bitonic_merge_zmm<vtype>(zmm_t2); + zmm[2] = bitonic_merge_zmm<vtype>(zmm_t3); + zmm[3] = bitonic_merge_zmm<vtype>(zmm_t4); + zmm[4] = bitonic_merge_zmm<vtype>(zmm_t5); + zmm[5] = bitonic_merge_zmm<vtype>(zmm_t6); + zmm[6] = bitonic_merge_zmm<vtype>(zmm_t7); + zmm[7] = bitonic_merge_zmm<vtype>(zmm_t8); +} + +template <typename vtype, typename type_t> +static NPY_INLINE void +sort_16(type_t *arr, npy_int N) +{ + __mmask16 load_mask = (0x0001 << N) - 0x0001; + typename vtype::zmm_t zmm = + vtype::mask_loadu(vtype::zmm_max(), load_mask, arr); + vtype::mask_storeu(arr, load_mask, sort_zmm<vtype>(zmm)); +} + +template <typename vtype, typename type_t> +static NPY_INLINE void +sort_32(type_t *arr, npy_int N) +{ + if (N <= 16) { + sort_16<vtype>(arr, N); + return; + } + using zmm_t = typename vtype::zmm_t; + zmm_t zmm1 = vtype::loadu(arr); + __mmask16 load_mask = (0x0001 << (N - 16)) - 0x0001; + zmm_t zmm2 = vtype::mask_loadu(vtype::zmm_max(), load_mask, arr + 16); + zmm1 = sort_zmm<vtype>(zmm1); + zmm2 = sort_zmm<vtype>(zmm2); + bitonic_merge_two_zmm<vtype>(&zmm1, &zmm2); + vtype::storeu(arr, zmm1); + vtype::mask_storeu(arr + 16, load_mask, zmm2); +} + +template <typename vtype, typename type_t> +static NPY_INLINE void +sort_64(type_t *arr, npy_int N) +{ + if (N <= 32) { + sort_32<vtype>(arr, N); + return; + } + using zmm_t = typename vtype::zmm_t; + zmm_t zmm[4]; + zmm[0] = vtype::loadu(arr); + zmm[1] = vtype::loadu(arr + 16); + __mmask16 load_mask1 = 0xFFFF, load_mask2 = 0xFFFF; + if (N < 48) { + load_mask1 = (0x0001 << (N - 32)) - 0x0001; + load_mask2 = 0x0000; + } + else if (N < 64) { + load_mask2 = (0x0001 << (N - 48)) - 0x0001; + } + zmm[2] = vtype::mask_loadu(vtype::zmm_max(), load_mask1, arr + 32); + zmm[3] = vtype::mask_loadu(vtype::zmm_max(), load_mask2, arr + 48); + zmm[0] = sort_zmm<vtype>(zmm[0]); + zmm[1] = sort_zmm<vtype>(zmm[1]); + zmm[2] = sort_zmm<vtype>(zmm[2]); + zmm[3] = sort_zmm<vtype>(zmm[3]); + bitonic_merge_two_zmm<vtype>(&zmm[0], &zmm[1]); + bitonic_merge_two_zmm<vtype>(&zmm[2], &zmm[3]); + bitonic_merge_four_zmm<vtype>(zmm); + vtype::storeu(arr, zmm[0]); + vtype::storeu(arr + 16, zmm[1]); + vtype::mask_storeu(arr + 32, load_mask1, zmm[2]); + vtype::mask_storeu(arr + 48, load_mask2, zmm[3]); +} + +template <typename vtype, typename type_t> +static NPY_INLINE void +sort_128(type_t *arr, npy_int N) +{ + if (N <= 64) { + sort_64<vtype>(arr, N); + return; + } + using zmm_t = typename vtype::zmm_t; + zmm_t zmm[8]; + zmm[0] = vtype::loadu(arr); + zmm[1] = vtype::loadu(arr + 16); + zmm[2] = vtype::loadu(arr + 32); + zmm[3] = vtype::loadu(arr + 48); + zmm[0] = sort_zmm<vtype>(zmm[0]); + zmm[1] = sort_zmm<vtype>(zmm[1]); + zmm[2] = sort_zmm<vtype>(zmm[2]); + zmm[3] = sort_zmm<vtype>(zmm[3]); + __mmask16 load_mask1 = 0xFFFF, load_mask2 = 0xFFFF; + __mmask16 load_mask3 = 0xFFFF, load_mask4 = 0xFFFF; + if (N < 80) { + load_mask1 = (0x0001 << (N - 64)) - 0x0001; + load_mask2 = 0x0000; + load_mask3 = 0x0000; + load_mask4 = 0x0000; + } + else if (N < 96) { + load_mask2 = (0x0001 << (N - 80)) - 0x0001; + load_mask3 = 0x0000; + load_mask4 = 0x0000; + } + else if (N < 112) { + load_mask3 = (0x0001 << (N - 96)) - 0x0001; + load_mask4 = 0x0000; + } + else { + load_mask4 = (0x0001 << (N - 112)) - 0x0001; + } + zmm[4] = vtype::mask_loadu(vtype::zmm_max(), load_mask1, arr + 64); + zmm[5] = vtype::mask_loadu(vtype::zmm_max(), load_mask2, arr + 80); + zmm[6] = vtype::mask_loadu(vtype::zmm_max(), load_mask3, arr + 96); + zmm[7] = vtype::mask_loadu(vtype::zmm_max(), load_mask4, arr + 112); + zmm[4] = sort_zmm<vtype>(zmm[4]); + zmm[5] = sort_zmm<vtype>(zmm[5]); + zmm[6] = sort_zmm<vtype>(zmm[6]); + zmm[7] = sort_zmm<vtype>(zmm[7]); + bitonic_merge_two_zmm<vtype>(&zmm[0], &zmm[1]); + bitonic_merge_two_zmm<vtype>(&zmm[2], &zmm[3]); + bitonic_merge_two_zmm<vtype>(&zmm[4], &zmm[5]); + bitonic_merge_two_zmm<vtype>(&zmm[6], &zmm[7]); + bitonic_merge_four_zmm<vtype>(zmm); + bitonic_merge_four_zmm<vtype>(zmm + 4); + bitonic_merge_eight_zmm<vtype>(zmm); + vtype::storeu(arr, zmm[0]); + vtype::storeu(arr + 16, zmm[1]); + vtype::storeu(arr + 32, zmm[2]); + vtype::storeu(arr + 48, zmm[3]); + vtype::mask_storeu(arr + 64, load_mask1, zmm[4]); + vtype::mask_storeu(arr + 80, load_mask2, zmm[5]); + vtype::mask_storeu(arr + 96, load_mask3, zmm[6]); + vtype::mask_storeu(arr + 112, load_mask4, zmm[7]); +} + +template <typename type_t> +static NPY_INLINE void +swap(type_t *arr, npy_intp ii, npy_intp jj) +{ + type_t temp = arr[ii]; + arr[ii] = arr[jj]; + arr[jj] = temp; +} + +// Median of 3 strategy +// template<typename type_t> +// static NPY_INLINE +// npy_intp get_pivot_index(type_t *arr, const npy_intp left, const npy_intp +// right) { +// return (rand() % (right + 1 - left)) + left; +// //npy_intp middle = ((right-left)/2) + left; +// //type_t a = arr[left], b = arr[middle], c = arr[right]; +// //if ((b >= a && b <= c) || (b <= a && b >= c)) +// // return middle; +// //if ((a >= b && a <= c) || (a <= b && a >= c)) +// // return left; +// //else +// // return right; +//} + +/* + * Picking the pivot: Median of 72 array elements chosen at random. + */ + +template <typename vtype, typename type_t> +static NPY_INLINE type_t +get_pivot(type_t *arr, const npy_intp left, const npy_intp right) +{ + /* seeds for vectorized random number generator */ + __m256i s0 = _mm256_setr_epi64x(8265987198341093849, 3762817312854612374, + 1324281658759788278, 6214952190349879213); + __m256i s1 = _mm256_setr_epi64x(2874178529384792648, 1257248936691237653, + 7874578921548791257, 1998265912745817298); + s0 = _mm256_add_epi64(s0, _mm256_set1_epi64x(left)); + s1 = _mm256_sub_epi64(s1, _mm256_set1_epi64x(right)); + + npy_intp arrsize = right - left + 1; + __m256i bound = + _mm256_set1_epi32(arrsize > INT32_MAX ? INT32_MAX : arrsize); + __m512i left_vec = _mm512_set1_epi64(left); + __m512i right_vec = _mm512_set1_epi64(right); + using ymm_t = typename vtype::ymm_t; + ymm_t v[9]; + /* fill 9 vectors with random numbers */ + for (npy_int i = 0; i < 9; ++i) { + __m256i rand_64 = vnext(&s0, &s1); /* vector with 4 random uint64_t */ + __m512i rand_32 = _mm512_cvtepi32_epi64(rnd_epu32( + rand_64, bound)); /* random numbers between 0 and bound - 1 */ + __m512i indices; + if (i < 5) + indices = + _mm512_add_epi64(left_vec, rand_32); /* indices for arr */ + else + indices = + _mm512_sub_epi64(right_vec, rand_32); /* indices for arr */ + + v[i] = vtype::template i64gather<sizeof(type_t)>(indices, arr); + } + + /* median network for 9 elements */ + COEX<vtype>(v[0], v[1]); + COEX<vtype>(v[2], v[3]); + COEX<vtype>(v[4], v[5]); + COEX<vtype>(v[6], v[7]); + COEX<vtype>(v[0], v[2]); + COEX<vtype>(v[1], v[3]); + COEX<vtype>(v[4], v[6]); + COEX<vtype>(v[5], v[7]); + COEX<vtype>(v[0], v[4]); + COEX<vtype>(v[1], v[2]); + COEX<vtype>(v[5], v[6]); + COEX<vtype>(v[3], v[7]); + COEX<vtype>(v[1], v[5]); + COEX<vtype>(v[2], v[6]); + COEX<vtype>(v[3], v[5]); + COEX<vtype>(v[2], v[4]); + COEX<vtype>(v[3], v[4]); + COEX<vtype>(v[3], v[8]); + COEX<vtype>(v[4], v[8]); + + // technically v[4] needs to be sorted before we pick the correct median, + // picking the 4th element works just as well for performance + type_t *temp = (type_t *)&v[4]; + + return temp[4]; +} + +/* + * Parition one ZMM register based on the pivot and returns the index of the + * last element that is less than equal to the pivot. + */ +template <typename vtype, typename type_t, typename zmm_t> +static NPY_INLINE npy_int +partition_vec(type_t *arr, npy_intp left, npy_intp right, const zmm_t curr_vec, + const zmm_t pivot_vec, zmm_t *smallest_vec, zmm_t *biggest_vec) +{ + /* which elements are larger than the pivot */ + __mmask16 gt_mask = vtype::ge(curr_vec, pivot_vec); + npy_int amount_gt_pivot = _mm_popcnt_u32((npy_int)gt_mask); + vtype::mask_compressstoreu(arr + left, _knot_mask16(gt_mask), curr_vec); + vtype::mask_compressstoreu(arr + right - amount_gt_pivot, gt_mask, + curr_vec); + *smallest_vec = vtype::min(curr_vec, *smallest_vec); + *biggest_vec = vtype::max(curr_vec, *biggest_vec); + return amount_gt_pivot; +} + +/* + * Parition an array based on the pivot and returns the index of the + * last element that is less than equal to the pivot. + */ +template <typename vtype, typename type_t> +static NPY_INLINE npy_intp +partition_avx512(type_t *arr, npy_intp left, npy_intp right, type_t pivot, + type_t *smallest, type_t *biggest) +{ + /* make array length divisible by 16 , shortening the array */ + for (npy_int i = (right - left) % 16; i > 0; --i) { + *smallest = MIN(*smallest, arr[left]); + *biggest = MAX(*biggest, arr[left]); + if (arr[left] > pivot) { + swap(arr, left, --right); + } + else { + ++left; + } + } + + if (left == right) + return left; /* less than 16 elements in the array */ + + using zmm_t = typename vtype::zmm_t; + zmm_t pivot_vec = vtype::set1(pivot); + zmm_t min_vec = vtype::set1(*smallest); + zmm_t max_vec = vtype::set1(*biggest); + + if (right - left == 16) { + zmm_t vec = vtype::loadu(arr + left); + npy_int amount_gt_pivot = partition_vec<vtype>( + arr, left, left + 16, vec, pivot_vec, &min_vec, &max_vec); + *smallest = vtype::reducemin(min_vec); + *biggest = vtype::reducemax(max_vec); + return left + (16 - amount_gt_pivot); + } + + // first and last 16 values are partitioned at the end + zmm_t vec_left = vtype::loadu(arr + left); + zmm_t vec_right = vtype::loadu(arr + (right - 16)); + // store points of the vectors + npy_intp r_store = right - 16; + npy_intp l_store = left; + // indices for loading the elements + left += 16; + right -= 16; + while (right - left != 0) { + zmm_t curr_vec; + /* + * if fewer elements are stored on the right side of the array, + * then next elements are loaded from the right side, + * otherwise from the left side + */ + if ((r_store + 16) - right < left - l_store) { + right -= 16; + curr_vec = vtype::loadu(arr + right); + } + else { + curr_vec = vtype::loadu(arr + left); + left += 16; + } + // partition the current vector and save it on both sides of the array + npy_int amount_gt_pivot = + partition_vec<vtype>(arr, l_store, r_store + 16, curr_vec, + pivot_vec, &min_vec, &max_vec); + ; + r_store -= amount_gt_pivot; + l_store += (16 - amount_gt_pivot); + } + + /* partition and save vec_left and vec_right */ + npy_int amount_gt_pivot = + partition_vec<vtype>(arr, l_store, r_store + 16, vec_left, + pivot_vec, &min_vec, &max_vec); + l_store += (16 - amount_gt_pivot); + amount_gt_pivot = + partition_vec<vtype>(arr, l_store, l_store + 16, vec_right, + pivot_vec, &min_vec, &max_vec); + l_store += (16 - amount_gt_pivot); + *smallest = vtype::reducemin(min_vec); + *biggest = vtype::reducemax(max_vec); + return l_store; +} + +template <typename vtype, typename type_t> +static NPY_INLINE void +qsort_(type_t *arr, npy_intp left, npy_intp right, npy_int max_iters) +{ + /* + * Resort to heapsort if quicksort isnt making any progress + */ + if (max_iters <= 0) { + heapsort_<typename vtype::tag>(arr + left, right + 1 - left); + return; + } + /* + * Base case: use bitonic networks to sort arrays <= 128 + */ + if (right + 1 - left <= 128) { + sort_128<vtype>(arr + left, right + 1 - left); + return; + } + + type_t pivot = get_pivot<vtype>(arr, left, right); + type_t smallest = vtype::type_max(); + type_t biggest = vtype::type_min(); + npy_intp pivot_index = partition_avx512<vtype>(arr, left, right + 1, pivot, + &smallest, &biggest); + if (pivot != smallest) + qsort_<vtype>(arr, left, pivot_index - 1, max_iters - 1); + if (pivot != biggest) + qsort_<vtype>(arr, pivot_index, right, max_iters - 1); +} + +static NPY_INLINE npy_intp +replace_nan_with_inf(npy_float *arr, npy_intp arrsize) +{ + npy_intp nan_count = 0; + __mmask16 loadmask = 0xFFFF; + while (arrsize > 0) { + if (arrsize < 16) { + loadmask = (0x0001 << arrsize) - 0x0001; + } + __m512 in_zmm = _mm512_maskz_loadu_ps(loadmask, arr); + __mmask16 nanmask = _mm512_cmp_ps_mask(in_zmm, in_zmm, _CMP_NEQ_UQ); + nan_count += _mm_popcnt_u32((npy_int)nanmask); + _mm512_mask_storeu_ps(arr, nanmask, ZMM_MAX_FLOAT); + arr += 16; + arrsize -= 16; + } + return nan_count; +} + +static NPY_INLINE void +replace_inf_with_nan(npy_float *arr, npy_intp arrsize, npy_intp nan_count) +{ + for (npy_intp ii = arrsize - 1; nan_count > 0; --ii) { + arr[ii] = NPY_NANF; + nan_count -= 1; + } +} + +/*************************************** + * C > C++ dispatch + ***************************************/ + +NPY_NO_EXPORT void +NPY_CPU_DISPATCH_CURFX(x86_quicksort_int)(void *arr, npy_intp arrsize) +{ + if (arrsize > 1) { + qsort_<vector<npy_int>, npy_int>((npy_int *)arr, 0, arrsize - 1, + 2 * log2(arrsize)); + } +} + +NPY_NO_EXPORT void +NPY_CPU_DISPATCH_CURFX(x86_quicksort_uint)(void *arr, npy_intp arrsize) +{ + if (arrsize > 1) { + qsort_<vector<npy_uint>, npy_uint>((npy_uint *)arr, 0, arrsize - 1, + 2 * log2(arrsize)); + } +} + +NPY_NO_EXPORT void +NPY_CPU_DISPATCH_CURFX(x86_quicksort_float)(void *arr, npy_intp arrsize) +{ + if (arrsize > 1) { + npy_intp nan_count = replace_nan_with_inf((npy_float *)arr, arrsize); + qsort_<vector<npy_float>, npy_float>((npy_float *)arr, 0, arrsize - 1, + 2 * log2(arrsize)); + replace_inf_with_nan((npy_float *)arr, arrsize, nan_count); + } +} + +#endif // NPY_HAVE_AVX512_SKX diff --git a/numpy/core/src/umath/fast_loop_macros.h b/numpy/core/src/umath/fast_loop_macros.h index 4a36c9721..6132d0d71 100644 --- a/numpy/core/src/umath/fast_loop_macros.h +++ b/numpy/core/src/umath/fast_loop_macros.h @@ -103,6 +103,9 @@ abs_ptrdiff(char *a, char *b) && (steps[0] == steps[2])\ && (steps[0] == 0)) +/* input contiguous (for binary reduces only) */ +#define IS_BINARY_REDUCE_INPUT_CONT(tin) (steps[1] == sizeof(tin)) + /* binary loop input and output contiguous */ #define IS_BINARY_CONT(tin, tout) (steps[0] == sizeof(tin) && \ steps[1] == sizeof(tin) && \ @@ -252,6 +255,34 @@ abs_ptrdiff(char *a, char *b) TYPE io1 = *(TYPE *)iop1; \ BINARY_REDUCE_LOOP_INNER +/* + * op should be the code working on `TYPE in2` and + * reading/storing the result in `TYPE *io1` + */ +#define BASE_BINARY_REDUCE_LOOP(TYPE, op) \ + BINARY_REDUCE_LOOP_INNER { \ + const TYPE in2 = *(TYPE *)ip2; \ + op; \ + } + +#define BINARY_REDUCE_LOOP_FAST_INNER(TYPE, op)\ + /* condition allows compiler to optimize the generic macro */ \ + if(IS_BINARY_REDUCE_INPUT_CONT(TYPE)) { \ + BASE_BINARY_REDUCE_LOOP(TYPE, op) \ + } \ + else { \ + BASE_BINARY_REDUCE_LOOP(TYPE, op) \ + } + +#define BINARY_REDUCE_LOOP_FAST(TYPE, op)\ + do { \ + char *iop1 = args[0]; \ + TYPE io1 = *(TYPE *)iop1; \ + BINARY_REDUCE_LOOP_FAST_INNER(TYPE, op); \ + *((TYPE *)iop1) = io1; \ + } \ + while (0) + #define IS_BINARY_STRIDE_ONE(esize, vsize) \ ((steps[0] == esize) && \ (steps[1] == esize) && \ diff --git a/numpy/core/src/umath/loops.c.src b/numpy/core/src/umath/loops.c.src index 7f084ac39..252eac481 100644 --- a/numpy/core/src/umath/loops.c.src +++ b/numpy/core/src/umath/loops.c.src @@ -636,10 +636,7 @@ NPY_NO_EXPORT NPY_GCC_OPT_3 @ATTR@ void @TYPE@_@kind@@isa@(char **args, npy_intp const *dimensions, npy_intp const *steps, void *NPY_UNUSED(func)) { if (IS_BINARY_REDUCE) { - BINARY_REDUCE_LOOP(@type@) { - io1 @OP@= *(@type@ *)ip2; - } - *((@type@ *)iop1) = io1; + BINARY_REDUCE_LOOP_FAST(@type@, io1 @OP@= in2); } else { BINARY_LOOP_FAST(@type@, @type@, *out = in1 @OP@ in2); diff --git a/numpy/core/tests/test_dlpack.py b/numpy/core/tests/test_dlpack.py index 43a663f66..717210b54 100644 --- a/numpy/core/tests/test_dlpack.py +++ b/numpy/core/tests/test_dlpack.py @@ -27,12 +27,12 @@ class TestDLPack: z = y['int'] with pytest.raises(RuntimeError): - np._from_dlpack(z) + np.from_dlpack(z) @pytest.mark.skipif(IS_PYPY, reason="PyPy can't get refcounts.") def test_from_dlpack_refcount(self): x = np.arange(5) - y = np._from_dlpack(x) + y = np.from_dlpack(x) assert sys.getrefcount(x) == 3 del y assert sys.getrefcount(x) == 2 @@ -45,7 +45,7 @@ class TestDLPack: ]) def test_dtype_passthrough(self, dtype): x = np.arange(5, dtype=dtype) - y = np._from_dlpack(x) + y = np.from_dlpack(x) assert y.dtype == x.dtype assert_array_equal(x, y) @@ -54,44 +54,44 @@ class TestDLPack: x = np.asarray(np.datetime64('2021-05-27')) with pytest.raises(TypeError): - np._from_dlpack(x) + np.from_dlpack(x) def test_invalid_byte_swapping(self): dt = np.dtype('=i8').newbyteorder() x = np.arange(5, dtype=dt) with pytest.raises(TypeError): - np._from_dlpack(x) + np.from_dlpack(x) def test_non_contiguous(self): x = np.arange(25).reshape((5, 5)) y1 = x[0] - assert_array_equal(y1, np._from_dlpack(y1)) + assert_array_equal(y1, np.from_dlpack(y1)) y2 = x[:, 0] - assert_array_equal(y2, np._from_dlpack(y2)) + assert_array_equal(y2, np.from_dlpack(y2)) y3 = x[1, :] - assert_array_equal(y3, np._from_dlpack(y3)) + assert_array_equal(y3, np.from_dlpack(y3)) y4 = x[1] - assert_array_equal(y4, np._from_dlpack(y4)) + assert_array_equal(y4, np.from_dlpack(y4)) y5 = np.diagonal(x).copy() - assert_array_equal(y5, np._from_dlpack(y5)) + assert_array_equal(y5, np.from_dlpack(y5)) @pytest.mark.parametrize("ndim", range(33)) def test_higher_dims(self, ndim): shape = (1,) * ndim x = np.zeros(shape, dtype=np.float64) - assert shape == np._from_dlpack(x).shape + assert shape == np.from_dlpack(x).shape def test_dlpack_device(self): x = np.arange(5) assert x.__dlpack_device__() == (1, 0) - y = np._from_dlpack(x) + y = np.from_dlpack(x) assert y.__dlpack_device__() == (1, 0) z = y[::2] assert z.__dlpack_device__() == (1, 0) @@ -113,11 +113,11 @@ class TestDLPack: def test_ndim0(self): x = np.array(1.0) - y = np._from_dlpack(x) + y = np.from_dlpack(x) assert_array_equal(x, y) def test_size1dims_arrays(self): x = np.ndarray(dtype='f8', shape=(10, 5, 1), strides=(8, 80, 4), buffer=np.ones(1000, dtype=np.uint8), order='F') - y = np._from_dlpack(x) + y = np.from_dlpack(x) assert_array_equal(x, y) diff --git a/numpy/core/tests/test_dtype.py b/numpy/core/tests/test_dtype.py index 1a8e747e1..a23840403 100644 --- a/numpy/core/tests/test_dtype.py +++ b/numpy/core/tests/test_dtype.py @@ -1381,6 +1381,11 @@ def test_keyword_argument(): assert np.dtype(dtype=np.float64) == np.dtype(np.float64) +def test_ulong_dtype(): + # test for gh-21063 + assert np.dtype("ulong") == np.dtype(np.uint) + + class TestFromDTypeAttribute: def test_simple(self): class dt: diff --git a/numpy/core/tests/test_numerictypes.py b/numpy/core/tests/test_numerictypes.py index 9cb00342d..73ff4764d 100644 --- a/numpy/core/tests/test_numerictypes.py +++ b/numpy/core/tests/test_numerictypes.py @@ -436,6 +436,13 @@ class TestSctypeDict: assert_(np.sctypeDict['f8'] is not np.longdouble) assert_(np.sctypeDict['c16'] is not np.clongdouble) + def test_ulong(self): + # Test that 'ulong' behaves like 'long'. np.sctypeDict['long'] is an + # alias for np.int_, but np.long is not supported for historical + # reasons (gh-21063) + assert_(np.sctypeDict['ulong'] is np.uint) + assert_(not hasattr(np, 'ulong')) + class TestBitName: def test_abstract(self): diff --git a/numpy/distutils/ccompiler_opt.py b/numpy/distutils/ccompiler_opt.py index d0c40a7b2..0343cb8f4 100644 --- a/numpy/distutils/ccompiler_opt.py +++ b/numpy/distutils/ccompiler_opt.py @@ -2551,6 +2551,8 @@ class CCompilerOpt(_Config, _Distutils, _Cache, _CCompiler, _Feature, _Parse): except OSError: pass + os.makedirs(os.path.dirname(config_path), exist_ok=True) + self.dist_log("generate dispatched config -> ", config_path) dispatch_calls = [] for tar in targets: diff --git a/numpy/f2py/capi_maps.py b/numpy/f2py/capi_maps.py index b4fab71f9..1f43207db 100644 --- a/numpy/f2py/capi_maps.py +++ b/numpy/f2py/capi_maps.py @@ -190,7 +190,7 @@ def load_f2cmap_file(f2cmap_file): try: outmess('Reading f2cmap from {!r} ...\n'.format(f2cmap_file)) with open(f2cmap_file, 'r') as f: - d = eval(f.read(), {}, {}) + d = eval(f.read().lower(), {}, {}) for k, d1 in d.items(): for k1 in d1.keys(): d1[k1.lower()] = d1[k1] @@ -505,7 +505,7 @@ def sign2map(a, var): init,init.r,init.i,pytype vardebuginfo,vardebugshowvalue,varshowvalue varrformat - + intent """ out_a = a diff --git a/numpy/lib/__init__.pyi b/numpy/lib/__init__.pyi index 7338fc6d6..0e3da5b41 100644 --- a/numpy/lib/__init__.pyi +++ b/numpy/lib/__init__.pyi @@ -14,7 +14,7 @@ from numpy.lib import ( format as format, mixins as mixins, scimath as scimath, - stride_tricks as stride_stricks, + stride_tricks as stride_tricks, ) from numpy.lib._version import ( diff --git a/numpy/linalg/linalg.py b/numpy/linalg/linalg.py index d3acc5938..36c5eb85c 100644 --- a/numpy/linalg/linalg.py +++ b/numpy/linalg/linalg.py @@ -899,11 +899,11 @@ def qr(a, mode='reduced'): [1, 1], [1, 1], [2, 1]]) - >>> b = np.array([1, 0, 2, 1]) + >>> b = np.array([1, 2, 2, 3]) >>> q, r = np.linalg.qr(A) >>> p = np.dot(q.T, b) >>> np.dot(np.linalg.inv(r), p) - array([ 1.1e-16, 1.0e+00]) + array([ 1., 1.]) """ if mode not in ('reduced', 'complete', 'r', 'raw'): diff --git a/numpy/random/_generator.pyx b/numpy/random/_generator.pyx index d7c1879e7..3eb5876ec 100644 --- a/numpy/random/_generator.pyx +++ b/numpy/random/_generator.pyx @@ -15,6 +15,7 @@ from numpy.core.multiarray import normalize_axis_index from .c_distributions cimport * from libc cimport string +from libc.math cimport sqrt from libc.stdint cimport (uint8_t, uint16_t, uint32_t, uint64_t, int32_t, int64_t, INT64_MAX, SIZE_MAX) from ._bounded_integers cimport (_rand_bool, _rand_int32, _rand_int64, @@ -2997,6 +2998,22 @@ cdef class Generator: then the probability distribution of the number of non-"1"s that appear before the third "1" is a negative binomial distribution. + Because this method internally calls ``Generator.poisson`` with an + intermediate random value, a ValueError is raised when the choice of + :math:`n` and :math:`p` would result in the mean + 10 sigma of the sampled + intermediate distribution exceeding the max acceptable value of the + ``Generator.poisson`` method. This happens when :math:`p` is too low + (a lot of failures happen for every success) and :math:`n` is too big ( + a lot of sucesses are allowed). + Therefore, the :math:`n` and :math:`p` values must satisfy the constraint: + + .. math:: n\\frac{1-p}{p}+10n\\sqrt{n}\\frac{1-p}{p}<2^{63}-1-10\\sqrt{2^{63}-1}, + + Where the left side of the equation is the derived mean + 10 sigma of + a sample from the gamma distribution internally used as the :math:`lam` + parameter of a poisson sample, and the right side of the equation is + the constraint for maximum value of :math:`lam` in ``Generator.poisson``. + References ---------- .. [1] Weisstein, Eric W. "Negative Binomial Distribution." From @@ -3021,9 +3038,41 @@ cdef class Generator: ... print(i, "wells drilled, probability of one success =", probability) """ + + cdef bint is_scalar = True + cdef double *_dn + cdef double *_dp + cdef double _dmax_lam + + p_arr = <np.ndarray>np.PyArray_FROM_OTF(p, np.NPY_DOUBLE, np.NPY_ALIGNED) + is_scalar = is_scalar and np.PyArray_NDIM(p_arr) == 0 + n_arr = <np.ndarray>np.PyArray_FROM_OTF(n, np.NPY_DOUBLE, np.NPY_ALIGNED) + is_scalar = is_scalar and np.PyArray_NDIM(n_arr) == 0 + + if not is_scalar: + check_array_constraint(n_arr, 'n', CONS_POSITIVE_NOT_NAN) + check_array_constraint(p_arr, 'p', CONS_BOUNDED_GT_0_1) + # Check that the choice of negative_binomial parameters won't result in a + # call to the poisson distribution function with a value of lam too large. + max_lam_arr = (1 - p_arr) / p_arr * (n_arr + 10 * np.sqrt(n_arr)) + if np.any(np.greater(max_lam_arr, POISSON_LAM_MAX)): + raise ValueError("n too large or p too small, see Generator.negative_binomial Notes") + + else: + _dn = <double*>np.PyArray_DATA(n_arr) + _dp = <double*>np.PyArray_DATA(p_arr) + + check_constraint(_dn[0], 'n', CONS_POSITIVE_NOT_NAN) + check_constraint(_dp[0], 'p', CONS_BOUNDED_GT_0_1) + # Check that the choice of negative_binomial parameters won't result in a + # call to the poisson distribution function with a value of lam too large. + _dmax_lam = (1 - _dp[0]) / _dp[0] * (_dn[0] + 10 * sqrt(_dn[0])) + if _dmax_lam > POISSON_LAM_MAX: + raise ValueError("n too large or p too small, see Generator.negative_binomial Notes") + return disc(&random_negative_binomial, &self._bitgen, size, self.lock, 2, 0, - n, 'n', CONS_POSITIVE_NOT_NAN, - p, 'p', CONS_BOUNDED_GT_0_1, + n_arr, 'n', CONS_NONE, + p_arr, 'p', CONS_NONE, 0.0, '', CONS_NONE) def poisson(self, lam=1.0, size=None): diff --git a/numpy/random/tests/test_generator_mt19937.py b/numpy/random/tests/test_generator_mt19937.py index 7c61038a4..3ccb9103c 100644 --- a/numpy/random/tests/test_generator_mt19937.py +++ b/numpy/random/tests/test_generator_mt19937.py @@ -1485,6 +1485,13 @@ class TestRandomDist: with assert_raises(ValueError): x = random.negative_binomial(1, 0) + def test_negative_binomial_invalid_p_n_combination(self): + # Verify that values of p and n that would result in an overflow + # or infinite loop raise an exception. + with np.errstate(invalid='ignore'): + assert_raises(ValueError, random.negative_binomial, 2**62, 0.1) + assert_raises(ValueError, random.negative_binomial, [2**62], [0.1]) + def test_noncentral_chisquare(self): random = Generator(MT19937(self.seed)) actual = random.noncentral_chisquare(df=5, nonc=5, size=(3, 2)) diff --git a/numpy/typing/_char_codes.py b/numpy/typing/_char_codes.py index 139471084..f840d17bb 100644 --- a/numpy/typing/_char_codes.py +++ b/numpy/typing/_char_codes.py @@ -30,7 +30,7 @@ _UByteCodes = Literal["ubyte", "B", "=B", "<B", ">B"] _UShortCodes = Literal["ushort", "H", "=H", "<H", ">H"] _UIntCCodes = Literal["uintc", "I", "=I", "<I", ">I"] _UIntPCodes = Literal["uintp", "uint0", "P", "=P", "<P", ">P"] -_UIntCodes = Literal["uint", "L", "=L", "<L", ">L"] +_UIntCodes = Literal["ulong", "uint", "L", "=L", "<L", ">L"] _ULongLongCodes = Literal["ulonglong", "Q", "=Q", "<Q", ">Q"] _HalfCodes = Literal["half", "e", "=e", "<e", ">e"] diff --git a/numpy/typing/tests/data/reveal/array_constructors.pyi b/numpy/typing/tests/data/reveal/array_constructors.pyi index 52aabd126..4a865a31c 100644 --- a/numpy/typing/tests/data/reveal/array_constructors.pyi +++ b/numpy/typing/tests/data/reveal/array_constructors.pyi @@ -168,7 +168,7 @@ reveal_type(np.full(1, i8, dtype=np.float64)) # E: ndarray[Any, dtype[{float64} reveal_type(np.full(1, i8, dtype=float)) # E: ndarray[Any, dtype[Any]] reveal_type(np.indices([1, 2, 3])) # E: ndarray[Any, dtype[{int_}]] -reveal_type(np.indices([1, 2, 3], sparse=True)) # E: tuple[ndarray[Any, dtype[{int_}]]] +reveal_type(np.indices([1, 2, 3], sparse=True)) # E: tuple[ndarray[Any, dtype[{int_}]], ...] reveal_type(np.fromfunction(func, (3, 5))) # E: SubClass[{float64}] diff --git a/numpy/typing/tests/data/reveal/arrayterator.pyi b/numpy/typing/tests/data/reveal/arrayterator.pyi index 2dab9d08c..b6c26ddb7 100644 --- a/numpy/typing/tests/data/reveal/arrayterator.pyi +++ b/numpy/typing/tests/data/reveal/arrayterator.pyi @@ -9,7 +9,7 @@ reveal_type(ar_iter.buf_size) # E: Union[None, builtins.int] reveal_type(ar_iter.start) # E: builtins.list[builtins.int] reveal_type(ar_iter.stop) # E: builtins.list[builtins.int] reveal_type(ar_iter.step) # E: builtins.list[builtins.int] -reveal_type(ar_iter.shape) # E: builtins.tuple[builtins.int] +reveal_type(ar_iter.shape) # E: builtins.tuple[builtins.int, ...] reveal_type(ar_iter.flat) # E: typing.Generator[{int64}, None, None] reveal_type(ar_iter.__array__()) # E: ndarray[Any, dtype[{int64}]] diff --git a/numpy/typing/tests/data/reveal/dtype.pyi b/numpy/typing/tests/data/reveal/dtype.pyi index 934d7da5e..7a658511a 100644 --- a/numpy/typing/tests/data/reveal/dtype.pyi +++ b/numpy/typing/tests/data/reveal/dtype.pyi @@ -48,11 +48,11 @@ reveal_type(np.dtype(("U", 10))) # E: dtype[void] # Methods and attributes reveal_type(dtype_U.base) # E: dtype[Any] -reveal_type(dtype_U.subdtype) # E: Union[None, Tuple[dtype[Any], builtins.tuple[builtins.int]]] +reveal_type(dtype_U.subdtype) # E: Union[None, Tuple[dtype[Any], builtins.tuple[builtins.int, ...]]] reveal_type(dtype_U.newbyteorder()) # E: dtype[str_] reveal_type(dtype_U.type) # E: Type[str_] reveal_type(dtype_U.name) # E: str -reveal_type(dtype_U.names) # E: Union[None, builtins.tuple[builtins.str]] +reveal_type(dtype_U.names) # E: Union[None, builtins.tuple[builtins.str, ...]] reveal_type(dtype_U * 0) # E: dtype[str_] reveal_type(dtype_U * 1) # E: dtype[str_] diff --git a/numpy/typing/tests/data/reveal/flatiter.pyi b/numpy/typing/tests/data/reveal/flatiter.pyi index 9e7d6e545..0f0758175 100644 --- a/numpy/typing/tests/data/reveal/flatiter.pyi +++ b/numpy/typing/tests/data/reveal/flatiter.pyi @@ -5,9 +5,9 @@ a: np.flatiter[np.ndarray[Any, np.dtype[np.str_]]] reveal_type(a.base) # E: ndarray[Any, dtype[str_]] reveal_type(a.copy()) # E: ndarray[Any, dtype[str_]] -reveal_type(a.coords) # E: tuple[builtins.int] +reveal_type(a.coords) # E: tuple[builtins.int, ...] reveal_type(a.index) # E: int -reveal_type(iter(a)) # E: Iterator[str_] +reveal_type(iter(a)) # E: flatiter[ndarray[Any, dtype[str_]]] reveal_type(next(a)) # E: str_ reveal_type(a[0]) # E: str_ reveal_type(a[[0, 1, 2]]) # E: ndarray[Any, dtype[str_]] diff --git a/numpy/typing/tests/data/reveal/fromnumeric.pyi b/numpy/typing/tests/data/reveal/fromnumeric.pyi index 6adbc35bf..e769abcf5 100644 --- a/numpy/typing/tests/data/reveal/fromnumeric.pyi +++ b/numpy/typing/tests/data/reveal/fromnumeric.pyi @@ -119,17 +119,17 @@ reveal_type(np.ravel(f)) # E: ndarray[Any, dtype[Any]] reveal_type(np.ravel(AR_b)) # E: ndarray[Any, dtype[bool_]] reveal_type(np.ravel(AR_f4)) # E: ndarray[Any, dtype[{float32}]] -reveal_type(np.nonzero(b)) # E: tuple[ndarray[Any, dtype[{intp}]]] -reveal_type(np.nonzero(f4)) # E: tuple[ndarray[Any, dtype[{intp}]]] -reveal_type(np.nonzero(f)) # E: tuple[ndarray[Any, dtype[{intp}]]] -reveal_type(np.nonzero(AR_b)) # E: tuple[ndarray[Any, dtype[{intp}]]] -reveal_type(np.nonzero(AR_f4)) # E: tuple[ndarray[Any, dtype[{intp}]]] - -reveal_type(np.shape(b)) # E: tuple[builtins.int] -reveal_type(np.shape(f4)) # E: tuple[builtins.int] -reveal_type(np.shape(f)) # E: tuple[builtins.int] -reveal_type(np.shape(AR_b)) # E: tuple[builtins.int] -reveal_type(np.shape(AR_f4)) # E: tuple[builtins.int] +reveal_type(np.nonzero(b)) # E: tuple[ndarray[Any, dtype[{intp}]], ...] +reveal_type(np.nonzero(f4)) # E: tuple[ndarray[Any, dtype[{intp}]], ...] +reveal_type(np.nonzero(f)) # E: tuple[ndarray[Any, dtype[{intp}]], ...] +reveal_type(np.nonzero(AR_b)) # E: tuple[ndarray[Any, dtype[{intp}]], ...] +reveal_type(np.nonzero(AR_f4)) # E: tuple[ndarray[Any, dtype[{intp}]], ...] + +reveal_type(np.shape(b)) # E: tuple[builtins.int, ...] +reveal_type(np.shape(f4)) # E: tuple[builtins.int, ...] +reveal_type(np.shape(f)) # E: tuple[builtins.int, ...] +reveal_type(np.shape(AR_b)) # E: tuple[builtins.int, ...] +reveal_type(np.shape(AR_f4)) # E: tuple[builtins.int, ...] reveal_type(np.compress([True], b)) # E: ndarray[Any, dtype[bool_]] reveal_type(np.compress([True], f4)) # E: ndarray[Any, dtype[{float32}]] diff --git a/numpy/typing/tests/data/reveal/index_tricks.pyi b/numpy/typing/tests/data/reveal/index_tricks.pyi index 4018605ea..707d6f3d4 100644 --- a/numpy/typing/tests/data/reveal/index_tricks.pyi +++ b/numpy/typing/tests/data/reveal/index_tricks.pyi @@ -16,22 +16,22 @@ reveal_type(np.ndenumerate(AR_i8).iter) # E: flatiter[ndarray[Any, dtype[{int64 reveal_type(np.ndenumerate(AR_LIKE_f).iter) # E: flatiter[ndarray[Any, dtype[{double}]]] reveal_type(np.ndenumerate(AR_LIKE_U).iter) # E: flatiter[ndarray[Any, dtype[str_]]] -reveal_type(next(np.ndenumerate(AR_i8))) # E: Tuple[builtins.tuple[builtins.int], {int64}] -reveal_type(next(np.ndenumerate(AR_LIKE_f))) # E: Tuple[builtins.tuple[builtins.int], {double}] -reveal_type(next(np.ndenumerate(AR_LIKE_U))) # E: Tuple[builtins.tuple[builtins.int], str_] +reveal_type(next(np.ndenumerate(AR_i8))) # E: Tuple[builtins.tuple[builtins.int, ...], {int64}] +reveal_type(next(np.ndenumerate(AR_LIKE_f))) # E: Tuple[builtins.tuple[builtins.int, ...], {double}] +reveal_type(next(np.ndenumerate(AR_LIKE_U))) # E: Tuple[builtins.tuple[builtins.int, ...], str_] -reveal_type(iter(np.ndenumerate(AR_i8))) # E: Iterator[Tuple[builtins.tuple[builtins.int], {int64}]] -reveal_type(iter(np.ndenumerate(AR_LIKE_f))) # E: Iterator[Tuple[builtins.tuple[builtins.int], {double}]] -reveal_type(iter(np.ndenumerate(AR_LIKE_U))) # E: Iterator[Tuple[builtins.tuple[builtins.int], str_]] +reveal_type(iter(np.ndenumerate(AR_i8))) # E: ndenumerate[{int64}] +reveal_type(iter(np.ndenumerate(AR_LIKE_f))) # E: ndenumerate[{double}] +reveal_type(iter(np.ndenumerate(AR_LIKE_U))) # E: ndenumerate[str_] reveal_type(np.ndindex(1, 2, 3)) # E: numpy.ndindex reveal_type(np.ndindex((1, 2, 3))) # E: numpy.ndindex -reveal_type(iter(np.ndindex(1, 2, 3))) # E: Iterator[builtins.tuple[builtins.int]] -reveal_type(next(np.ndindex(1, 2, 3))) # E: builtins.tuple[builtins.int] +reveal_type(iter(np.ndindex(1, 2, 3))) # E: ndindex +reveal_type(next(np.ndindex(1, 2, 3))) # E: builtins.tuple[builtins.int, ...] -reveal_type(np.unravel_index([22, 41, 37], (7, 6))) # E: tuple[ndarray[Any, dtype[{intp}]]] -reveal_type(np.unravel_index([31, 41, 13], (7, 6), order="F")) # E: tuple[ndarray[Any, dtype[{intp}]]] -reveal_type(np.unravel_index(1621, (6, 7, 8, 9))) # E: tuple[{intp}] +reveal_type(np.unravel_index([22, 41, 37], (7, 6))) # E: tuple[ndarray[Any, dtype[{intp}]], ...] +reveal_type(np.unravel_index([31, 41, 13], (7, 6), order="F")) # E: tuple[ndarray[Any, dtype[{intp}]], ...] +reveal_type(np.unravel_index(1621, (6, 7, 8, 9))) # E: tuple[{intp}, ...] reveal_type(np.ravel_multi_index([[1]], (7, 6))) # E: ndarray[Any, dtype[{intp}]] reveal_type(np.ravel_multi_index(AR_LIKE_i, (7, 6))) # E: {intp} @@ -54,13 +54,13 @@ reveal_type(np.s_[0:1]) # E: builtins.slice reveal_type(np.s_[0:1, None:3]) # E: Tuple[builtins.slice, builtins.slice] reveal_type(np.s_[0, 0:1, ..., [0, 1, 3]]) # E: Tuple[Literal[0]?, builtins.slice, builtins.ellipsis, builtins.list[builtins.int]] -reveal_type(np.ix_(AR_LIKE_b)) # E: tuple[ndarray[Any, dtype[bool_]]] -reveal_type(np.ix_(AR_LIKE_i, AR_LIKE_f)) # E: tuple[ndarray[Any, dtype[{double}]]] -reveal_type(np.ix_(AR_i8)) # E: tuple[ndarray[Any, dtype[{int64}]]] +reveal_type(np.ix_(AR_LIKE_b)) # E: tuple[ndarray[Any, dtype[bool_]], ...] +reveal_type(np.ix_(AR_LIKE_i, AR_LIKE_f)) # E: tuple[ndarray[Any, dtype[{double}]], ...] +reveal_type(np.ix_(AR_i8)) # E: tuple[ndarray[Any, dtype[{int64}]], ...] reveal_type(np.fill_diagonal(AR_i8, 5)) # E: None -reveal_type(np.diag_indices(4)) # E: tuple[ndarray[Any, dtype[{int_}]]] -reveal_type(np.diag_indices(2, 3)) # E: tuple[ndarray[Any, dtype[{int_}]]] +reveal_type(np.diag_indices(4)) # E: tuple[ndarray[Any, dtype[{int_}]], ...] +reveal_type(np.diag_indices(2, 3)) # E: tuple[ndarray[Any, dtype[{int_}]], ...] -reveal_type(np.diag_indices_from(AR_i8)) # E: tuple[ndarray[Any, dtype[{int_}]]] +reveal_type(np.diag_indices_from(AR_i8)) # E: tuple[ndarray[Any, dtype[{int_}]], ...] diff --git a/numpy/typing/tests/data/reveal/multiarray.pyi b/numpy/typing/tests/data/reveal/multiarray.pyi index f8937c379..27a54f50d 100644 --- a/numpy/typing/tests/data/reveal/multiarray.pyi +++ b/numpy/typing/tests/data/reveal/multiarray.pyi @@ -34,29 +34,29 @@ timedelta_seq: list[dt.timedelta] def func(a: int) -> bool: ... -reveal_type(next(b_f8)) # E: tuple[Any] +reveal_type(next(b_f8)) # E: tuple[Any, ...] reveal_type(b_f8.reset()) # E: None reveal_type(b_f8.index) # E: int -reveal_type(b_f8.iters) # E: tuple[flatiter[Any]] +reveal_type(b_f8.iters) # E: tuple[flatiter[Any], ...] reveal_type(b_f8.nd) # E: int reveal_type(b_f8.ndim) # E: int reveal_type(b_f8.numiter) # E: int -reveal_type(b_f8.shape) # E: tuple[builtins.int] +reveal_type(b_f8.shape) # E: tuple[builtins.int, ...] reveal_type(b_f8.size) # E: int -reveal_type(next(b_i8_f8_f8)) # E: tuple[Any] +reveal_type(next(b_i8_f8_f8)) # E: tuple[Any, ...] reveal_type(b_i8_f8_f8.reset()) # E: None reveal_type(b_i8_f8_f8.index) # E: int -reveal_type(b_i8_f8_f8.iters) # E: tuple[flatiter[Any]] +reveal_type(b_i8_f8_f8.iters) # E: tuple[flatiter[Any], ...] reveal_type(b_i8_f8_f8.nd) # E: int reveal_type(b_i8_f8_f8.ndim) # E: int reveal_type(b_i8_f8_f8.numiter) # E: int -reveal_type(b_i8_f8_f8.shape) # E: tuple[builtins.int] +reveal_type(b_i8_f8_f8.shape) # E: tuple[builtins.int, ...] reveal_type(b_i8_f8_f8.size) # E: int reveal_type(np.inner(AR_f8, AR_i8)) # E: Any -reveal_type(np.where([True, True, False])) # E: tuple[ndarray[Any, dtype[{intp}]]] +reveal_type(np.where([True, True, False])) # E: tuple[ndarray[Any, dtype[{intp}]], ...] reveal_type(np.where([True, True, False], 1, 0)) # E: ndarray[Any, dtype[Any]] reveal_type(np.lexsort([0, 1, 2])) # E: Any @@ -138,7 +138,7 @@ reveal_type(np.compare_chararrays(b"a", b"a", "==", True)) # E: ndarray[Any, dt reveal_type(np.add_docstring(func, "test")) # E: None -reveal_type(np.nested_iters([AR_i8, AR_i8], [[0], [1]], flags=["c_index"])) # E: tuple[nditer] -reveal_type(np.nested_iters([AR_i8, AR_i8], [[0], [1]], op_flags=[["readonly", "readonly"]])) # E: tuple[nditer] -reveal_type(np.nested_iters([AR_i8, AR_i8], [[0], [1]], op_dtypes=np.int_)) # E: tuple[nditer] -reveal_type(np.nested_iters([AR_i8, AR_i8], [[0], [1]], order="C", casting="no")) # E: tuple[nditer] +reveal_type(np.nested_iters([AR_i8, AR_i8], [[0], [1]], flags=["c_index"])) # E: tuple[nditer, ...] +reveal_type(np.nested_iters([AR_i8, AR_i8], [[0], [1]], op_flags=[["readonly", "readonly"]])) # E: tuple[nditer, ...] +reveal_type(np.nested_iters([AR_i8, AR_i8], [[0], [1]], op_dtypes=np.int_)) # E: tuple[nditer, ...] +reveal_type(np.nested_iters([AR_i8, AR_i8], [[0], [1]], order="C", casting="no")) # E: tuple[nditer, ...] diff --git a/numpy/typing/tests/data/reveal/ndarray_misc.pyi b/numpy/typing/tests/data/reveal/ndarray_misc.pyi index c9a42b3e8..c990e2ab1 100644 --- a/numpy/typing/tests/data/reveal/ndarray_misc.pyi +++ b/numpy/typing/tests/data/reveal/ndarray_misc.pyi @@ -165,7 +165,7 @@ reveal_type(AR_f8.dot(1)) # E: ndarray[Any, Any] reveal_type(AR_f8.dot([1])) # E: Any reveal_type(AR_f8.dot(1, out=B)) # E: SubClass -reveal_type(AR_f8.nonzero()) # E: tuple[ndarray[Any, dtype[{intp}]]] +reveal_type(AR_f8.nonzero()) # E: tuple[ndarray[Any, dtype[{intp}]], ...] reveal_type(AR_f8.searchsorted(1)) # E: {intp} reveal_type(AR_f8.searchsorted([1])) # E: ndarray[Any, dtype[{intp}]] diff --git a/numpy/typing/tests/data/reveal/nditer.pyi b/numpy/typing/tests/data/reveal/nditer.pyi index 65861da54..fd8b7e109 100644 --- a/numpy/typing/tests/data/reveal/nditer.pyi +++ b/numpy/typing/tests/data/reveal/nditer.pyi @@ -7,7 +7,7 @@ reveal_type(np.nditer([0, 1], op_flags=[["readonly", "readonly"]])) # E: nditer reveal_type(np.nditer([0, 1], op_dtypes=np.int_)) # E: nditer reveal_type(np.nditer([0, 1], order="C", casting="no")) # E: nditer -reveal_type(nditer_obj.dtypes) # E: tuple[dtype[Any]] +reveal_type(nditer_obj.dtypes) # E: tuple[dtype[Any], ...] reveal_type(nditer_obj.finished) # E: bool reveal_type(nditer_obj.has_delayed_bufalloc) # E: bool reveal_type(nditer_obj.has_index) # E: bool @@ -15,15 +15,15 @@ reveal_type(nditer_obj.has_multi_index) # E: bool reveal_type(nditer_obj.index) # E: int reveal_type(nditer_obj.iterationneedsapi) # E: bool reveal_type(nditer_obj.iterindex) # E: int -reveal_type(nditer_obj.iterrange) # E: tuple[builtins.int] +reveal_type(nditer_obj.iterrange) # E: tuple[builtins.int, ...] reveal_type(nditer_obj.itersize) # E: int -reveal_type(nditer_obj.itviews) # E: tuple[ndarray[Any, dtype[Any]]] -reveal_type(nditer_obj.multi_index) # E: tuple[builtins.int] +reveal_type(nditer_obj.itviews) # E: tuple[ndarray[Any, dtype[Any]], ...] +reveal_type(nditer_obj.multi_index) # E: tuple[builtins.int, ...] reveal_type(nditer_obj.ndim) # E: int reveal_type(nditer_obj.nop) # E: int -reveal_type(nditer_obj.operands) # E: tuple[ndarray[Any, dtype[Any]]] -reveal_type(nditer_obj.shape) # E: tuple[builtins.int] -reveal_type(nditer_obj.value) # E: tuple[ndarray[Any, dtype[Any]]] +reveal_type(nditer_obj.operands) # E: tuple[ndarray[Any, dtype[Any]], ...] +reveal_type(nditer_obj.shape) # E: tuple[builtins.int, ...] +reveal_type(nditer_obj.value) # E: tuple[ndarray[Any, dtype[Any]], ...] reveal_type(nditer_obj.close()) # E: None reveal_type(nditer_obj.copy()) # E: nditer @@ -35,12 +35,12 @@ reveal_type(nditer_obj.remove_multi_index()) # E: None reveal_type(nditer_obj.reset()) # E: None reveal_type(len(nditer_obj)) # E: int -reveal_type(iter(nditer_obj)) # E: Iterator[builtins.tuple[ndarray[Any, dtype[Any]]]] -reveal_type(next(nditer_obj)) # E: tuple[ndarray[Any, dtype[Any]]] +reveal_type(iter(nditer_obj)) # E: nditer +reveal_type(next(nditer_obj)) # E: tuple[ndarray[Any, dtype[Any]], ...] reveal_type(nditer_obj.__copy__()) # E: nditer with nditer_obj as f: reveal_type(f) # E: nditer reveal_type(nditer_obj[0]) # E: ndarray[Any, dtype[Any]] -reveal_type(nditer_obj[:]) # E: tuple[ndarray[Any, dtype[Any]]] +reveal_type(nditer_obj[:]) # E: tuple[ndarray[Any, dtype[Any]], ...] nditer_obj[0] = 0 nditer_obj[:] = [0, 1] diff --git a/numpy/typing/tests/data/reveal/numeric.pyi b/numpy/typing/tests/data/reveal/numeric.pyi index 246d79be8..b8fe15d3a 100644 --- a/numpy/typing/tests/data/reveal/numeric.pyi +++ b/numpy/typing/tests/data/reveal/numeric.pyi @@ -105,11 +105,11 @@ reveal_type(np.cross(AR_i8, AR_c16)) # E: ndarray[Any, dtype[complexfloating[An reveal_type(np.cross(AR_O, AR_O)) # E: ndarray[Any, dtype[object_]] reveal_type(np.indices([0, 1, 2])) # E: ndarray[Any, dtype[{int_}]] -reveal_type(np.indices([0, 1, 2], sparse=True)) # E: tuple[ndarray[Any, dtype[{int_}]]] +reveal_type(np.indices([0, 1, 2], sparse=True)) # E: tuple[ndarray[Any, dtype[{int_}]], ...] reveal_type(np.indices([0, 1, 2], dtype=np.float64)) # E: ndarray[Any, dtype[{float64}]] -reveal_type(np.indices([0, 1, 2], sparse=True, dtype=np.float64)) # E: tuple[ndarray[Any, dtype[{float64}]]] +reveal_type(np.indices([0, 1, 2], sparse=True, dtype=np.float64)) # E: tuple[ndarray[Any, dtype[{float64}]], ...] reveal_type(np.indices([0, 1, 2], dtype=float)) # E: ndarray[Any, dtype[Any]] -reveal_type(np.indices([0, 1, 2], sparse=True, dtype=float)) # E: tuple[ndarray[Any, dtype[Any]]] +reveal_type(np.indices([0, 1, 2], sparse=True, dtype=float)) # E: tuple[ndarray[Any, dtype[Any]], ...] reveal_type(np.binary_repr(1)) # E: str diff --git a/numpy/typing/tests/data/reveal/stride_tricks.pyi b/numpy/typing/tests/data/reveal/stride_tricks.pyi index 4eeb42095..17769dc4b 100644 --- a/numpy/typing/tests/data/reveal/stride_tricks.pyi +++ b/numpy/typing/tests/data/reveal/stride_tricks.pyi @@ -21,8 +21,8 @@ reveal_type(np.broadcast_to(AR_f8, 5)) # E: ndarray[Any, dtype[{float64}]] reveal_type(np.broadcast_to(AR_LIKE_f, (1, 5))) # E: ndarray[Any, dtype[Any]] reveal_type(np.broadcast_to(AR_f8, [4, 6], subok=True)) # E: ndarray[Any, dtype[{float64}]] -reveal_type(np.broadcast_shapes((1, 2), [3, 1], (3, 2))) # E: tuple[builtins.int] -reveal_type(np.broadcast_shapes((6, 7), (5, 6, 1), 7, (5, 1, 7))) # E: tuple[builtins.int] +reveal_type(np.broadcast_shapes((1, 2), [3, 1], (3, 2))) # E: tuple[builtins.int, ...] +reveal_type(np.broadcast_shapes((6, 7), (5, 6, 1), 7, (5, 1, 7))) # E: tuple[builtins.int, ...] reveal_type(np.broadcast_arrays(AR_f8, AR_f8)) # E: list[ndarray[Any, dtype[Any]]] reveal_type(np.broadcast_arrays(AR_f8, AR_LIKE_f)) # E: list[ndarray[Any, dtype[Any]]] diff --git a/numpy/typing/tests/data/reveal/testing.pyi b/numpy/typing/tests/data/reveal/testing.pyi index edd4bb3bf..5440af800 100644 --- a/numpy/typing/tests/data/reveal/testing.pyi +++ b/numpy/typing/tests/data/reveal/testing.pyi @@ -30,7 +30,7 @@ reveal_type(np.testing.clear_and_catch_warnings(modules=[np.testing])) # E: _cl reveal_type(np.testing.clear_and_catch_warnings(True)) # E: _clear_and_catch_warnings_with_records reveal_type(np.testing.clear_and_catch_warnings(False)) # E: _clear_and_catch_warnings_without_records reveal_type(np.testing.clear_and_catch_warnings(bool_obj)) # E: clear_and_catch_warnings -reveal_type(np.testing.clear_and_catch_warnings.class_modules) # E: tuple[types.ModuleType] +reveal_type(np.testing.clear_and_catch_warnings.class_modules) # E: tuple[types.ModuleType, ...] reveal_type(np.testing.clear_and_catch_warnings.modules) # E: set[types.ModuleType] with np.testing.clear_and_catch_warnings(True) as c1: diff --git a/test_requirements.txt b/test_requirements.txt index 7ca578e84..446708058 100644 --- a/test_requirements.txt +++ b/test_requirements.txt @@ -9,4 +9,5 @@ pytest-cov==3.0.0 cffi; python_version < '3.10' # For testing types. Notes on the restrictions: # - Mypy relies on C API features not present in PyPy -mypy==0.931; platform_python_implementation != "PyPy" +# NOTE: Keep mypy in sync with environment.yml +mypy==0.940; platform_python_implementation != "PyPy" diff --git a/tools/gitpod/settings.json b/tools/gitpod/settings.json index 8f070c04c..ea0775f68 100644 --- a/tools/gitpod/settings.json +++ b/tools/gitpod/settings.json @@ -5,5 +5,5 @@ "restructuredtext.updateOnTextChanged": "true", "restructuredtext.updateDelay": 300, "restructuredtext.linter.disabled": true, - "python.pythonPath": "/home/gitpod/mambaforge3/envs/numpy-dev/bin/python" + "python.defaultInterpreterPath": "/home/gitpod/mambaforge3/envs/numpy-dev/bin/python" }
\ No newline at end of file diff --git a/tools/refguide_check.py b/tools/refguide_check.py index 619d6c644..eb9a27ab5 100644 --- a/tools/refguide_check.py +++ b/tools/refguide_check.py @@ -104,6 +104,8 @@ DOCTEST_SKIPDICT = { 'numpy.random.vonmises': None, 'numpy.random.power': None, 'numpy.random.zipf': None, + # cases where NumPy docstrings import things from other 3'rd party libs: + 'numpy.core.from_dlpack': None, # remote / local file IO with DataSource is problematic in doctest: 'numpy.lib.DataSource': None, 'numpy.lib.Repository': None, diff --git a/tools/wheels/upload_wheels.sh b/tools/wheels/upload_wheels.sh index 130062d5e..21c6e629e 100644 --- a/tools/wheels/upload_wheels.sh +++ b/tools/wheels/upload_wheels.sh @@ -41,8 +41,16 @@ upload_wheels() { echo no token set, not uploading else python -m pip install git+https://github.com/Anaconda-Server/anaconda-client - ls ./wheelhouse/*.whl - anaconda -t ${TOKEN} upload --skip -u ${ANACONDA_ORG} ./wheelhouse/*.whl + if compgen -G "./wheelhouse/*.gz"; then + echo "Found sdist" + anaconda -t ${TOKEN} upload --skip -u ${ANACONDA_ORG} ./wheelhouse/*.gz + elif compgen -G "./wheelhouse/*.whl"; then + echo "Found wheel" + anaconda -t ${TOKEN} upload --skip -u ${ANACONDA_ORG} ./wheelhouse/*.whl + else + echo "Files do not exist" + return 1 + fi echo "PyPI-style index: https://pypi.anaconda.org/$ANACONDA_ORG/simple" fi fi |