Skip to content

Commit a739f1f

Browse files
ndgrigorianndgrigorian
committed
Update umath patching
* Add safety checks to patching * Add warnings related to multi-threaded programs * Implement _GlobalPatch wrapper class which implements patching globally * Rename _patch module to _patch_numpy * Rename patching functions
1 parent 6cc1b0b commit a739f1f

12 files changed

Lines changed: 520 additions & 128 deletions

File tree

‎.github/workflows/conda-package-cf.yml‎

Lines changed: 353 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,353 @@
1+
name: Conda package with conda-forge channel only
2+
3+
on:
4+
push:
5+
branches:
6+
- master
7+
pull_request:
8+
9+
permissions: read-all
10+
11+
env:
12+
PACKAGE_NAME: mkl_umath
13+
MODULE_NAME: mkl_umath
14+
TEST_ENV_NAME: test_mkl_umath
15+
VER_SCRIPT1: "import json; f = open('ver.json', 'r'); j = json.load(f); f.close(); d = j['mkl_umath'][0];"
16+
VER_SCRIPT2: "print('='.join((d[s] for s in ('version', 'build'))))"
17+
18+
jobs:
19+
build_linux:
20+
runs-on: ubuntu-latest
21+
strategy:
22+
matrix:
23+
include:
24+
- python: "3.10"
25+
numpy: "2.2"
26+
- python: "3.11"
27+
numpy: "2.3"
28+
- python: "3.12"
29+
numpy: "2.3"
30+
- python: "3.13"
31+
numpy: "2.3"
32+
- python: "3.14"
33+
numpy: "2.3"
34+
35+
steps:
36+
- name: Cancel Previous Runs
37+
uses: styfle/cancel-workflow-action@3155a141048f8f89c06b4cdae32e7853e97536bc # 0.13.0
38+
with:
39+
access_token: ${{ github.token }}
40+
- uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
41+
with:
42+
fetch-depth: 0
43+
44+
- name: Set pkgs_dirs
45+
run: |
46+
echo "pkgs_dirs: [~/.conda/pkgs]" >> ~/.condarc
47+
48+
- name: Cache conda packages
49+
uses: actions/cache@cdf6c1fa76f9f475f3d7449005a359c84ca0f306 # v5.0.3
50+
env:
51+
CACHE_NUMBER: 0 # Increase to reset cache
52+
with:
53+
path: ~/.conda/pkgs
54+
key:
55+
${{ runner.os }}-conda-${{ env.CACHE_NUMBER }}-python-${{ matrix.python }}-${{hashFiles('**/meta.yaml') }}
56+
restore-keys: |
57+
${{ runner.os }}-conda-${{ env.CACHE_NUMBER }}-python-${{ matrix.python }}-
58+
${{ runner.os }}-conda-${{ env.CACHE_NUMBER }}-
59+
60+
- name: Add conda to system path
61+
run: echo "$CONDA/bin" >> "$GITHUB_PATH"
62+
63+
- name: Install conda-build
64+
run: conda install conda-build
65+
66+
- name: Build conda package with NumPy 2.x
67+
run: |
68+
CHANNELS=(-c conda-forge --override-channels)
69+
VERSIONS=(--python "${{ matrix.python }}" --numpy "${{ matrix.numpy }}")
70+
TEST=(--no-test)
71+
echo "CONDA_BLD=${CONDA}/conda-bld/linux-64" >> "$GITHUB_ENV"
72+
73+
conda build \
74+
"${TEST[@]}" \
75+
"${VERSIONS[@]}" \
76+
"${CHANNELS[@]}" \
77+
conda-recipe-cf
78+
79+
- name: Upload artifact
80+
uses: actions/upload-artifact@bbbca2ddaa5d8feaa63e36b76fdaad77386f024f # v7.0.0
81+
with:
82+
name: ${{ env.PACKAGE_NAME }} ${{ runner.os }} Python ${{ matrix.python }}
83+
path: ${{ env.CONDA_BLD }}/${{ env.PACKAGE_NAME }}-*.conda
84+
85+
test_linux:
86+
needs: build_linux
87+
runs-on: ${{ matrix.runner }}
88+
89+
strategy:
90+
matrix:
91+
python: ["3.10", "3.11", "3.12", "3.13", "3.14"]
92+
numpy: ['numpy">2"']
93+
experimental: [false]
94+
runner: [ubuntu-latest]
95+
continue-on-error: ${{ matrix.experimental }}
96+
env:
97+
CHANNELS: -c conda-forge --override-channels
98+
99+
steps:
100+
- name: Download artifact
101+
uses: actions/download-artifact@70fc10c6e5e1ce46ad2ea6f2b72d43f7d47b13c3 # v8.0.0
102+
with:
103+
name: ${{ env.PACKAGE_NAME }} ${{ runner.os }} Python ${{ matrix.python }}
104+
- name: Add conda to system path
105+
run: echo "$CONDA/bin" >> "$GITHUB_PATH"
106+
- name: Install conda-build
107+
run: conda install conda-build
108+
- name: Create conda channel
109+
run: |
110+
mkdir -p "$GITHUB_WORKSPACE/channel/linux-64"
111+
mv "${PACKAGE_NAME}"-*.conda "$GITHUB_WORKSPACE/channel/linux-64"
112+
conda index "$GITHUB_WORKSPACE/channel"
113+
# Test channel
114+
conda search "$PACKAGE_NAME" -c "$GITHUB_WORKSPACE/channel" --override-channels
115+
116+
- name: Collect dependencies
117+
run: |
118+
CHANNELS=(-c "$GITHUB_WORKSPACE/channel" -c "conda-forge" --override-channels)
119+
conda create -n "${{ env.TEST_ENV_NAME }}" "$PACKAGE_NAME" "python=${{ matrix.python }}" "${CHANNELS[@]}" --only-deps --dry-run > lockfile
120+
- name: Display lockfile
121+
run: cat lockfile
122+
123+
- name: Set pkgs_dirs
124+
run: |
125+
echo "pkgs_dirs: [~/.conda/pkgs]" >> ~/.condarc
126+
127+
- name: Cache conda packages
128+
uses: actions/cache@cdf6c1fa76f9f475f3d7449005a359c84ca0f306 # v5.0.3
129+
env:
130+
CACHE_NUMBER: 0 # Increase to reset cache
131+
with:
132+
path: ~/.conda/pkgs
133+
key:
134+
${{ runner.os }}-conda-${{ env.CACHE_NUMBER }}-python-${{ matrix.python }}-${{hashFiles('lockfile') }}
135+
restore-keys: |
136+
${{ runner.os }}-conda-${{ env.CACHE_NUMBER }}-python-${{ matrix.python }}-
137+
${{ runner.os }}-conda-${{ env.CACHE_NUMBER }}-
138+
139+
- name: Install mkl_umath
140+
run: |
141+
CHANNELS=(-c "$GITHUB_WORKSPACE/channel" -c "conda-forge" --override-channels)
142+
conda create -n "${{ env.TEST_ENV_NAME }}" "python=${{ matrix.python }}" "$PACKAGE_NAME" pytest "${CHANNELS[@]}"
143+
# Test installed packages
144+
conda list -n "${{ env.TEST_ENV_NAME }}"
145+
146+
- name: Smoke test
147+
run: |
148+
source "$CONDA/etc/profile.d/conda.sh"
149+
conda activate "${{ env.TEST_ENV_NAME }}"
150+
python -c "import mkl_umath, numpy as np; mkl_umath.patch_numpy_umath(); np.sin(np.linspace(0, 1, num=10**6));"
151+
152+
- name: Run tests
153+
run: |
154+
source "$CONDA/etc/profile.d/conda.sh"
155+
conda activate "${{ env.TEST_ENV_NAME }}"
156+
pytest -v --pyargs ${{ env.MODULE_NAME }}
157+
158+
build_windows:
159+
runs-on: windows-latest
160+
defaults:
161+
run:
162+
shell: cmd /C CALL {0}
163+
strategy:
164+
matrix:
165+
include:
166+
- python: "3.10"
167+
numpy: "2.2"
168+
- python: "3.11"
169+
numpy: "2.3"
170+
- python: "3.12"
171+
numpy: "2.3"
172+
- python: "3.13"
173+
numpy: "2.3"
174+
- python: "3.14"
175+
numpy: "2.3"
176+
177+
env:
178+
conda-bld: C:\Miniconda\conda-bld\win-64\
179+
steps:
180+
- name: Cancel Previous Runs
181+
uses: styfle/cancel-workflow-action@3155a141048f8f89c06b4cdae32e7853e97536bc # 0.13.0
182+
with:
183+
access_token: ${{ github.token }}
184+
- uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
185+
with:
186+
fetch-depth: 0
187+
188+
- uses: conda-incubator/setup-miniconda@fc2d68f6413eb2d87b895e92f8584b5b94a10167 # v3.3.0
189+
with:
190+
miniforge-version: latest
191+
activate-environment: build
192+
channels: conda-forge
193+
python-version: ${{ matrix.python }}
194+
195+
- name: Cache conda packages
196+
uses: actions/cache@cdf6c1fa76f9f475f3d7449005a359c84ca0f306 # v5.0.3
197+
env:
198+
CACHE_NUMBER: 3 # Increase to reset cache
199+
with:
200+
path: /home/runner/conda_pkgs_dir
201+
key:
202+
${{ runner.os }}-conda-${{ env.CACHE_NUMBER }}-python-${{ matrix.python }}-${{hashFiles('**/meta.yaml') }}
203+
restore-keys: |
204+
${{ runner.os }}-conda-${{ env.CACHE_NUMBER }}-python-${{ matrix.python }}-
205+
${{ runner.os }}-conda-${{ env.CACHE_NUMBER }}-
206+
207+
- name: Store conda paths as envs
208+
shell: bash -l {0}
209+
run: |
210+
echo "CONDA_BLD=$CONDA/conda-bld/win-64/" | tr "\\\\" '/' >> "$GITHUB_ENV"
211+
212+
- name: Install conda build
213+
run: |
214+
conda install -n base -y conda-build
215+
conda list -n base
216+
217+
- name: Build conda package
218+
run: |
219+
conda build --no-test --python ${{ matrix.python }} --numpy ${{ matrix.numpy }} -c conda-forge --override-channels conda-recipe-cf
220+
221+
- name: Upload artifact
222+
uses: actions/upload-artifact@bbbca2ddaa5d8feaa63e36b76fdaad77386f024f # v7.0.0
223+
with:
224+
name: ${{ env.PACKAGE_NAME }} ${{ runner.os }} Python ${{ matrix.python }}
225+
path: ${{ env.CONDA_BLD }}${{ env.PACKAGE_NAME }}-*.conda
226+
227+
test_windows:
228+
needs: build_windows
229+
runs-on: ${{ matrix.runner }}
230+
defaults:
231+
run:
232+
shell: cmd /C CALL {0}
233+
strategy:
234+
matrix:
235+
python: ["3.10", "3.11", "3.12", "3.13", "3.14"]
236+
numpy: ['numpy">2"']
237+
experimental: [false]
238+
runner: [windows-latest]
239+
continue-on-error: ${{ matrix.experimental }}
240+
env:
241+
workdir: '${{ github.workspace }}'
242+
CHANNELS: -c conda-forge --override-channels
243+
244+
steps:
245+
- name: Download artifact
246+
uses: actions/download-artifact@70fc10c6e5e1ce46ad2ea6f2b72d43f7d47b13c3 # v8.0.0
247+
with:
248+
name: ${{ env.PACKAGE_NAME }} ${{ runner.os }} Python ${{ matrix.python }}
249+
250+
- uses: conda-incubator/setup-miniconda@fc2d68f6413eb2d87b895e92f8584b5b94a10167 # v3.3.0
251+
with:
252+
miniforge-version: latest
253+
channels: conda-forge
254+
activate-environment: ${{ env.TEST_ENV_NAME }}
255+
python-version: ${{ matrix.python }}
256+
257+
- name: Install conda-index
258+
run: |
259+
conda install -n base conda-index
260+
261+
- name: Create conda channel with the artifact bit
262+
shell: cmd /C CALL {0}
263+
run: |
264+
echo ${{ env.workdir }}
265+
mkdir ${{ env.workdir }}\channel
266+
mkdir ${{ env.workdir }}\channel\win-64
267+
move ${{ env.PACKAGE_NAME }}-*.conda ${{ env.workdir }}\channel\win-64
268+
dir ${{ env.workdir }}\channel\win-64
269+
270+
- name: Index the channel
271+
shell: cmd /C CALL {0}
272+
run: |
273+
conda index ${{ env.workdir }}\channel
274+
275+
- name: Dump mkl_umath version info from created channel to STDOUT
276+
shell: cmd /C CALL {0}
277+
run: |
278+
conda search ${{ env.PACKAGE_NAME }} -c ${{ env.workdir }}/channel --override-channels --info --json
279+
280+
- name: Dump mkl_umath version info from created channel into ver.json
281+
shell: cmd /C CALL {0}
282+
run: |
283+
conda search ${{ env.PACKAGE_NAME }} -c ${{ env.workdir }}/channel --override-channels --info --json > ${{ env.workdir }}\ver.json
284+
285+
- name: Output content of workdir
286+
shell: pwsh
287+
run: Get-ChildItem -Path ${{ env.workdir }}
288+
289+
- name: Output content of produced ver.json
290+
shell: pwsh
291+
run: Get-Content -Path ${{ env.workdir }}\ver.json
292+
293+
- name: Collect dependencies
294+
shell: cmd /C CALL {0}
295+
run: |
296+
@ECHO ON
297+
IF NOT EXIST ver.json (
298+
copy /Y ${{ env.workdir }}\ver.json .
299+
)
300+
SET "SCRIPT=%VER_SCRIPT1% %VER_SCRIPT2%"
301+
FOR /F "tokens=* USEBACKQ" %%F IN (`python -c "%SCRIPT%"`) DO (
302+
SET PACKAGE_VERSION=%%F
303+
)
304+
conda install -n ${{ env.TEST_ENV_NAME }} ${{ env.PACKAGE_NAME }}=%PACKAGE_VERSION% python=${{ matrix.python }} -c ${{ env.workdir }}/channel ${{ env.CHANNELS }} --only-deps --dry-run > lockfile
305+
306+
- name: Display lockfile content
307+
shell: pwsh
308+
run: Get-Content -Path .\lockfile
309+
310+
- name: Cache conda packages
311+
uses: actions/cache@cdf6c1fa76f9f475f3d7449005a359c84ca0f306 # v5.0.3
312+
env:
313+
CACHE_NUMBER: 0 # Increase to reset cache
314+
with:
315+
path: /home/runner/conda_pkgs_dir
316+
key:
317+
${{ runner.os }}-conda-${{ env.CACHE_NUMBER }}-python-${{ matrix.python }}-${{hashFiles('lockfile') }}
318+
restore-keys: |
319+
${{ runner.os }}-conda-${{ env.CACHE_NUMBER }}-python-${{ matrix.python }}-
320+
${{ runner.os }}-conda-${{ env.CACHE_NUMBER }}-
321+
322+
- name: Install mkl_umath
323+
shell: cmd /C CALL {0}
324+
run: |
325+
@ECHO ON
326+
IF NOT EXIST ver.json (
327+
copy /Y ${{ env.workdir }}\ver.json .
328+
)
329+
set "SCRIPT=%VER_SCRIPT1% %VER_SCRIPT2%"
330+
FOR /F "tokens=* USEBACKQ" %%F IN (`python -c "%SCRIPT%"`) DO (
331+
SET PACKAGE_VERSION=%%F
332+
)
333+
conda install -n ${{ env.TEST_ENV_NAME }} ${{ env.PACKAGE_NAME }}=%PACKAGE_VERSION% pytest python=${{ matrix.python }} -c ${{ env.workdir }}/channel ${{ env.CHANNELS }}
334+
335+
- name: Report content of test environment
336+
shell: cmd /C CALL {0}
337+
run: |
338+
echo "Value of CONDA environment variable was: " %CONDA%
339+
echo "Value of CONDA_PREFIX environment variable was: " %CONDA_PREFIX%
340+
conda info && conda list -n ${{ env.TEST_ENV_NAME }}
341+
342+
- name: Smoke test
343+
shell: cmd /C CALL {0}
344+
run: |
345+
@ECHO ON
346+
conda activate ${{ env.TEST_ENV_NAME }}
347+
python -c "import mkl_umath, numpy as np; mkl_umath.patch_numpy_umath(); np.sin(np.linspace(0, 1, num=10**6));"
348+
349+
- name: Run tests
350+
shell: cmd /C CALL {0}
351+
run: |
352+
conda activate ${{ env.TEST_ENV_NAME }}
353+
python -m pytest -v -s --pyargs ${{ env.MODULE_NAME }}

‎.github/workflows/conda-package.yml‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -133,7 +133,7 @@ jobs:
133133
run: |
134134
source "$CONDA/etc/profile.d/conda.sh"
135135
conda activate test_mkl_umath
136-
python -c "import mkl_umath, numpy as np; mkl_umath.use_in_numpy(); np.sin(np.linspace(0, 1, num=10**6));"
136+
python -c "import mkl_umath, numpy as np; mkl_umath.patch_numpy_umath(); np.sin(np.linspace(0, 1, num=10**6));"
137137
138138
- name: Run tests
139139
run: |
@@ -328,7 +328,7 @@ jobs:
328328
run: |
329329
@ECHO ON
330330
conda activate mkl_umath_test
331-
python -c "import mkl_umath, numpy as np; mkl_umath.use_in_numpy(); np.sin(np.linspace(0, 1, num=10**6));"
331+
python -c "import mkl_umath, numpy as np; mkl_umath.patch_numpy_umath(); np.sin(np.linspace(0, 1, num=10**6));"
332332
333333
- name: Run tests
334334
shell: cmd /C CALL {0}

‎AGENTS.md‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -7,13 +7,13 @@ Entry point for agent context in this repo.
77

88
It provides:
99
- `mkl_umath._ufuncs` — OneMKL-backed NumPy ufunc loops
10-
- `mkl_umath._patch` — runtime patching interface (`use_in_numpy()`, `restore()`, `is_patched()`)
10+
- `mkl_umath._patch_numpy` — runtime patching interface (`patch_numpy_umath` `restore_numpy_umath`, `is_patched()`)
1111
- Performance-optimized math operations (sin, cos, exp, log, etc.) using Intel MKL VM
1212

1313
## Key components
1414
- **Python interface:** `mkl_umath/__init__.py`, `_init_helper.py`
1515
- **Core C implementation:** `mkl_umath/src/` (ufuncsmodule.c, mkl_umath_loops.c.src)
16-
- **Cython patch layer:** `mkl_umath/src/_patch.pyx`
16+
- **Cython patch layer:** `mkl_umath/src/_patch_numpy.pyx`
1717
- **Code generation:** `generate_umath.py`, `generate_umath_doc.py`
1818
- **Build system:** CMake (CMakeLists.txt) + scikit-build
1919

@@ -50,9 +50,9 @@ CC=icx pip install --no-build-isolation --no-deps .
5050
## Usage
5151
```python
5252
import mkl_umath
53-
mkl_umath.use_in_numpy() # Patch NumPy to use MKL loops
53+
mkl_umath.patch_numpy_umath() # Patch NumPy to use MKL loops
5454
# ... perform NumPy operations (now accelerated) ...
55-
mkl_umath.restore() # Restore original NumPy loops
55+
mkl_umath.restore_numpy_umath() # Restore original NumPy loops
5656
```
5757

5858
## How to work in this repo

0 commit comments

Comments
 (0)