rarg-numba-patterns 0.0.1__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- rarg_numba_patterns-0.0.1/.github/dependabot.yml +6 -0
- rarg_numba_patterns-0.0.1/.github/workflows/ci.yml +30 -0
- rarg_numba_patterns-0.0.1/.gitignore +220 -0
- rarg_numba_patterns-0.0.1/.pre-commit-config.yaml +26 -0
- rarg_numba_patterns-0.0.1/LICENSE +28 -0
- rarg_numba_patterns-0.0.1/PKG-INFO +17 -0
- rarg_numba_patterns-0.0.1/README.md +3 -0
- rarg_numba_patterns-0.0.1/pyproject.toml +58 -0
- rarg_numba_patterns-0.0.1/src/rarg_numba_patterns/__init__.py +47 -0
- rarg_numba_patterns-0.0.1/src/rarg_numba_patterns/concurrency/__init__.py +0 -0
- rarg_numba_patterns-0.0.1/src/rarg_numba_patterns/concurrency/mutex.py +199 -0
- rarg_numba_patterns-0.0.1/src/rarg_numba_patterns/intrinsics.py +229 -0
- rarg_numba_patterns-0.0.1/src/rarg_numba_patterns/literals.py +215 -0
- rarg_numba_patterns-0.0.1/tests/conftest.py +14 -0
- rarg_numba_patterns-0.0.1/tests/test_intrinsics.py +33 -0
- rarg_numba_patterns-0.0.1/tests/test_literals.py +269 -0
- rarg_numba_patterns-0.0.1/tests/test_pointers.py +34 -0
- rarg_numba_patterns-0.0.1/tests/test_smoke.py +5 -0
|
@@ -0,0 +1,30 @@
|
|
|
1
|
+
name: Continuous Integration
|
|
2
|
+
|
|
3
|
+
on:
|
|
4
|
+
push: { branches: [main], tags: ["*"] }
|
|
5
|
+
pull_request:
|
|
6
|
+
schedule: [{ cron: '30 2 * * 1,4' }] # Every Monday and Thursday @ 2h30am UTC
|
|
7
|
+
|
|
8
|
+
jobs:
|
|
9
|
+
test:
|
|
10
|
+
uses: ratt-ru/rarg-gh-workflows/.github/workflows/python-ci.yml@main
|
|
11
|
+
with:
|
|
12
|
+
os-matrix: '["ubuntu-24.04"]' # Add other OS (macos-15) if necessary
|
|
13
|
+
install-command: 'pip install --group test .'
|
|
14
|
+
|
|
15
|
+
deploy:
|
|
16
|
+
needs: ["test"]
|
|
17
|
+
runs-on: ubuntu-latest
|
|
18
|
+
environment:
|
|
19
|
+
name: pypi
|
|
20
|
+
url: https://pypi.org/p/rarg-numba-patterns
|
|
21
|
+
permissions:
|
|
22
|
+
id-token: write
|
|
23
|
+
steps:
|
|
24
|
+
- uses: actions/setup-python@v6
|
|
25
|
+
with: { python-version: '3.13' }
|
|
26
|
+
- uses: actions/checkout@v6
|
|
27
|
+
- run: pip install -U pip build && python -m build --sdist --wheel
|
|
28
|
+
- name: Publish distribution 📦 to PyPI
|
|
29
|
+
if: github.event_name == 'push' && startsWith(github.event.ref, 'refs/tags')
|
|
30
|
+
uses: pypa/gh-action-pypi-publish@release/v1
|
|
@@ -0,0 +1,220 @@
|
|
|
1
|
+
# Byte-compiled / optimized / DLL files
|
|
2
|
+
__pycache__/
|
|
3
|
+
*.py[codz]
|
|
4
|
+
*$py.class
|
|
5
|
+
|
|
6
|
+
# C extensions
|
|
7
|
+
*.so
|
|
8
|
+
|
|
9
|
+
# Distribution / packaging
|
|
10
|
+
.Python
|
|
11
|
+
build/
|
|
12
|
+
develop-eggs/
|
|
13
|
+
dist/
|
|
14
|
+
downloads/
|
|
15
|
+
eggs/
|
|
16
|
+
.eggs/
|
|
17
|
+
lib/
|
|
18
|
+
lib64/
|
|
19
|
+
parts/
|
|
20
|
+
sdist/
|
|
21
|
+
var/
|
|
22
|
+
wheels/
|
|
23
|
+
share/python-wheels/
|
|
24
|
+
*.egg-info/
|
|
25
|
+
.installed.cfg
|
|
26
|
+
*.egg
|
|
27
|
+
MANIFEST
|
|
28
|
+
|
|
29
|
+
# PyInstaller
|
|
30
|
+
# Usually these files are written by a python script from a template
|
|
31
|
+
# before PyInstaller builds the exe, so as to inject date/other infos into it.
|
|
32
|
+
*.manifest
|
|
33
|
+
*.spec
|
|
34
|
+
|
|
35
|
+
# Installer logs
|
|
36
|
+
pip-log.txt
|
|
37
|
+
pip-delete-this-directory.txt
|
|
38
|
+
|
|
39
|
+
# Unit test / coverage reports
|
|
40
|
+
htmlcov/
|
|
41
|
+
.tox/
|
|
42
|
+
.nox/
|
|
43
|
+
.coverage
|
|
44
|
+
.coverage.*
|
|
45
|
+
.cache
|
|
46
|
+
nosetests.xml
|
|
47
|
+
coverage.xml
|
|
48
|
+
*.cover
|
|
49
|
+
*.py.cover
|
|
50
|
+
*.lcov
|
|
51
|
+
.hypothesis/
|
|
52
|
+
.pytest_cache/
|
|
53
|
+
cover/
|
|
54
|
+
|
|
55
|
+
# Translations
|
|
56
|
+
*.mo
|
|
57
|
+
*.pot
|
|
58
|
+
|
|
59
|
+
# Django stuff:
|
|
60
|
+
*.log
|
|
61
|
+
local_settings.py
|
|
62
|
+
db.sqlite3
|
|
63
|
+
db.sqlite3-journal
|
|
64
|
+
|
|
65
|
+
# Flask stuff:
|
|
66
|
+
instance/
|
|
67
|
+
.webassets-cache
|
|
68
|
+
|
|
69
|
+
# Scrapy stuff:
|
|
70
|
+
.scrapy
|
|
71
|
+
|
|
72
|
+
# Sphinx documentation
|
|
73
|
+
docs/_build/
|
|
74
|
+
|
|
75
|
+
# PyBuilder
|
|
76
|
+
.pybuilder/
|
|
77
|
+
target/
|
|
78
|
+
|
|
79
|
+
# Jupyter Notebook
|
|
80
|
+
.ipynb_checkpoints
|
|
81
|
+
|
|
82
|
+
# IPython
|
|
83
|
+
profile_default/
|
|
84
|
+
ipython_config.py
|
|
85
|
+
|
|
86
|
+
# pyenv
|
|
87
|
+
# For a library or package, you might want to ignore these files since the code is
|
|
88
|
+
# intended to run in multiple environments; otherwise, check them in:
|
|
89
|
+
# .python-version
|
|
90
|
+
|
|
91
|
+
# pipenv
|
|
92
|
+
# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
|
|
93
|
+
# However, in case of collaboration, if having platform-specific dependencies or dependencies
|
|
94
|
+
# having no cross-platform support, pipenv may install dependencies that don't work, or not
|
|
95
|
+
# install all needed dependencies.
|
|
96
|
+
# Pipfile.lock
|
|
97
|
+
|
|
98
|
+
# UV
|
|
99
|
+
# Similar to Pipfile.lock, it is generally recommended to include uv.lock in version control.
|
|
100
|
+
# This is especially recommended for binary packages to ensure reproducibility, and is more
|
|
101
|
+
# commonly ignored for libraries.
|
|
102
|
+
# uv.lock
|
|
103
|
+
|
|
104
|
+
# poetry
|
|
105
|
+
# Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
|
|
106
|
+
# This is especially recommended for binary packages to ensure reproducibility, and is more
|
|
107
|
+
# commonly ignored for libraries.
|
|
108
|
+
# https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
|
|
109
|
+
# poetry.lock
|
|
110
|
+
# poetry.toml
|
|
111
|
+
|
|
112
|
+
# pdm
|
|
113
|
+
# Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
|
|
114
|
+
# pdm recommends including project-wide configuration in pdm.toml, but excluding .pdm-python.
|
|
115
|
+
# https://pdm-project.org/en/latest/usage/project/#working-with-version-control
|
|
116
|
+
# pdm.lock
|
|
117
|
+
# pdm.toml
|
|
118
|
+
.pdm-python
|
|
119
|
+
.pdm-build/
|
|
120
|
+
|
|
121
|
+
# pixi
|
|
122
|
+
# Similar to Pipfile.lock, it is generally recommended to include pixi.lock in version control.
|
|
123
|
+
# pixi.lock
|
|
124
|
+
# Pixi creates a virtual environment in the .pixi directory, just like venv module creates one
|
|
125
|
+
# in the .venv directory. It is recommended not to include this directory in version control.
|
|
126
|
+
.pixi/*
|
|
127
|
+
!.pixi/config.toml
|
|
128
|
+
|
|
129
|
+
# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
|
|
130
|
+
__pypackages__/
|
|
131
|
+
|
|
132
|
+
# Celery stuff
|
|
133
|
+
celerybeat-schedule*
|
|
134
|
+
celerybeat.pid
|
|
135
|
+
|
|
136
|
+
# Redis
|
|
137
|
+
*.rdb
|
|
138
|
+
*.aof
|
|
139
|
+
*.pid
|
|
140
|
+
|
|
141
|
+
# RabbitMQ
|
|
142
|
+
mnesia/
|
|
143
|
+
rabbitmq/
|
|
144
|
+
rabbitmq-data/
|
|
145
|
+
|
|
146
|
+
# ActiveMQ
|
|
147
|
+
activemq-data/
|
|
148
|
+
|
|
149
|
+
# SageMath parsed files
|
|
150
|
+
*.sage.py
|
|
151
|
+
|
|
152
|
+
# Environments
|
|
153
|
+
.env
|
|
154
|
+
.envrc
|
|
155
|
+
.venv
|
|
156
|
+
env/
|
|
157
|
+
venv/
|
|
158
|
+
ENV/
|
|
159
|
+
env.bak/
|
|
160
|
+
venv.bak/
|
|
161
|
+
|
|
162
|
+
# Spyder project settings
|
|
163
|
+
.spyderproject
|
|
164
|
+
.spyproject
|
|
165
|
+
|
|
166
|
+
# Rope project settings
|
|
167
|
+
.ropeproject
|
|
168
|
+
|
|
169
|
+
# mkdocs documentation
|
|
170
|
+
/site
|
|
171
|
+
|
|
172
|
+
# mypy
|
|
173
|
+
.mypy_cache/
|
|
174
|
+
.dmypy.json
|
|
175
|
+
dmypy.json
|
|
176
|
+
|
|
177
|
+
# Pyre type checker
|
|
178
|
+
.pyre/
|
|
179
|
+
|
|
180
|
+
# pytype static type analyzer
|
|
181
|
+
.pytype/
|
|
182
|
+
|
|
183
|
+
# Cython debug symbols
|
|
184
|
+
cython_debug/
|
|
185
|
+
|
|
186
|
+
# PyCharm
|
|
187
|
+
# JetBrains specific template is maintained in a separate JetBrains.gitignore that can
|
|
188
|
+
# be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
|
|
189
|
+
# and can be added to the global gitignore or merged into this file. For a more nuclear
|
|
190
|
+
# option (not recommended) you can uncomment the following to ignore the entire idea folder.
|
|
191
|
+
# .idea/
|
|
192
|
+
|
|
193
|
+
# Abstra
|
|
194
|
+
# Abstra is an AI-powered process automation framework.
|
|
195
|
+
# Ignore directories containing user credentials, local state, and settings.
|
|
196
|
+
# Learn more at https://abstra.io/docs
|
|
197
|
+
.abstra/
|
|
198
|
+
|
|
199
|
+
# Visual Studio Code
|
|
200
|
+
# Visual Studio Code specific template is maintained in a separate VisualStudioCode.gitignore
|
|
201
|
+
# that can be found at https://github.com/github/gitignore/blob/main/Global/VisualStudioCode.gitignore
|
|
202
|
+
# and can be added to the global gitignore or merged into this file. However, if you prefer,
|
|
203
|
+
# you could uncomment the following to ignore the entire vscode folder
|
|
204
|
+
# .vscode/
|
|
205
|
+
# Temporary file for partial code execution
|
|
206
|
+
tempCodeRunnerFile.py
|
|
207
|
+
|
|
208
|
+
# Ruff stuff:
|
|
209
|
+
.ruff_cache/
|
|
210
|
+
|
|
211
|
+
# PyPI configuration file
|
|
212
|
+
.pypirc
|
|
213
|
+
|
|
214
|
+
# Marimo
|
|
215
|
+
marimo/_static/
|
|
216
|
+
marimo/_lsp/
|
|
217
|
+
__marimo__/
|
|
218
|
+
|
|
219
|
+
# Streamlit
|
|
220
|
+
.streamlit/secrets.toml
|
|
@@ -0,0 +1,26 @@
|
|
|
1
|
+
# See https://pre-commit.com for more information
|
|
2
|
+
# See https://pre-commit.com/hooks.html for more hooks
|
|
3
|
+
# Run `pre-commit autoupdate` after scaffolding to pin current hook revs.
|
|
4
|
+
repos:
|
|
5
|
+
- repo: https://github.com/pre-commit/pre-commit-hooks
|
|
6
|
+
rev: v5.0.0
|
|
7
|
+
hooks:
|
|
8
|
+
- id: trailing-whitespace
|
|
9
|
+
- id: end-of-file-fixer
|
|
10
|
+
- id: check-yaml
|
|
11
|
+
- id: check-added-large-files
|
|
12
|
+
- repo: https://github.com/astral-sh/ruff-pre-commit
|
|
13
|
+
rev: v0.12.3
|
|
14
|
+
hooks:
|
|
15
|
+
- id: ruff
|
|
16
|
+
# Handle imports, see
|
|
17
|
+
# https://github.com/astral-sh/ruff/issues/8232
|
|
18
|
+
# https://github.com/astral-sh/ruff/issues/10882
|
|
19
|
+
args: [ --fix, --extend-select, I ]
|
|
20
|
+
- id: ruff-format
|
|
21
|
+
name: ruff format
|
|
22
|
+
- repo: https://github.com/pre-commit/mirrors-mypy
|
|
23
|
+
rev: v1.17.0
|
|
24
|
+
hooks:
|
|
25
|
+
- id: mypy
|
|
26
|
+
name: mypy
|
|
@@ -0,0 +1,28 @@
|
|
|
1
|
+
BSD 3-Clause License
|
|
2
|
+
|
|
3
|
+
Copyright (c) 2026, Rhodes University Centre for Radio Astronomy Techniques & Technologies (RATT)
|
|
4
|
+
|
|
5
|
+
Redistribution and use in source and binary forms, with or without
|
|
6
|
+
modification, are permitted provided that the following conditions are met:
|
|
7
|
+
|
|
8
|
+
1. Redistributions of source code must retain the above copyright notice, this
|
|
9
|
+
list of conditions and the following disclaimer.
|
|
10
|
+
|
|
11
|
+
2. Redistributions in binary form must reproduce the above copyright notice,
|
|
12
|
+
this list of conditions and the following disclaimer in the documentation
|
|
13
|
+
and/or other materials provided with the distribution.
|
|
14
|
+
|
|
15
|
+
3. Neither the name of the copyright holder nor the names of its
|
|
16
|
+
contributors may be used to endorse or promote products derived from
|
|
17
|
+
this software without specific prior written permission.
|
|
18
|
+
|
|
19
|
+
THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
|
|
20
|
+
AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
|
|
21
|
+
IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
|
|
22
|
+
DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
|
|
23
|
+
FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
|
|
24
|
+
DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
|
|
25
|
+
SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
|
|
26
|
+
CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
|
|
27
|
+
OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
|
|
28
|
+
OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
|
@@ -0,0 +1,17 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: rarg-numba-patterns
|
|
3
|
+
Version: 0.0.1
|
|
4
|
+
Summary: General-purpose numba intrinsics and literals patterns
|
|
5
|
+
Project-URL: Homepage, https://github.com/ratt-ru/rarg-numba-patterns
|
|
6
|
+
Project-URL: Repository, https://github.com/ratt-ru/rarg-numba-patterns
|
|
7
|
+
Author-email: Simon Perkins <simon.perkins@gmail.com>
|
|
8
|
+
License-Expression: BSD-3-Clause
|
|
9
|
+
License-File: LICENSE
|
|
10
|
+
Requires-Python: <3.15,>=3.10
|
|
11
|
+
Requires-Dist: numba>=0.60.0
|
|
12
|
+
Requires-Dist: numpy>=2.0.0
|
|
13
|
+
Description-Content-Type: text/markdown
|
|
14
|
+
|
|
15
|
+
# rarg-numba-patterns
|
|
16
|
+
|
|
17
|
+
General-purpose numba intrinsics and literals patterns
|
|
@@ -0,0 +1,58 @@
|
|
|
1
|
+
[build-system]
|
|
2
|
+
requires = ["hatchling"]
|
|
3
|
+
build-backend = "hatchling.build"
|
|
4
|
+
|
|
5
|
+
[project]
|
|
6
|
+
name = "rarg-numba-patterns"
|
|
7
|
+
version = "0.0.1"
|
|
8
|
+
description = "General-purpose numba intrinsics and literals patterns"
|
|
9
|
+
readme = "README.md"
|
|
10
|
+
requires-python = ">=3.10, <3.15"
|
|
11
|
+
license = "BSD-3-Clause"
|
|
12
|
+
authors = [{ name = "Simon Perkins", email = "simon.perkins@gmail.com" }]
|
|
13
|
+
dependencies = [
|
|
14
|
+
"numpy>=2.0.0",
|
|
15
|
+
"numba>=0.60.0",
|
|
16
|
+
]
|
|
17
|
+
|
|
18
|
+
[project.urls]
|
|
19
|
+
Homepage = "https://github.com/ratt-ru/rarg-numba-patterns"
|
|
20
|
+
Repository = "https://github.com/ratt-ru/rarg-numba-patterns"
|
|
21
|
+
|
|
22
|
+
[dependency-groups]
|
|
23
|
+
test = ["pytest"]
|
|
24
|
+
dev = ["pre-commit", "mypy", "ruff", "tbump"]
|
|
25
|
+
|
|
26
|
+
[tool.hatch.build.targets.wheel]
|
|
27
|
+
packages = ["src/rarg_numba_patterns"]
|
|
28
|
+
|
|
29
|
+
[tool.ruff]
|
|
30
|
+
line-length = 88
|
|
31
|
+
indent-width = 2
|
|
32
|
+
|
|
33
|
+
[tool.ruff.lint]
|
|
34
|
+
extend-select = ["I"] # import sorting
|
|
35
|
+
|
|
36
|
+
[tool.mypy]
|
|
37
|
+
python_version = "3.10"
|
|
38
|
+
strict = true
|
|
39
|
+
|
|
40
|
+
[tool.pytest.ini_options]
|
|
41
|
+
testpaths = ["tests"]
|
|
42
|
+
|
|
43
|
+
[tool.tbump.version]
|
|
44
|
+
current = "0.0.1"
|
|
45
|
+
# https://semver.org/#spec-item-9
|
|
46
|
+
regex = '''(?P<major>\d+)\.(?P<minor>\d+)\.(?P<patch>\d+)-?(?P<ptype>(alpha|beta|rc)\.(?P<prerelease>\d+))?'''
|
|
47
|
+
|
|
48
|
+
[tool.tbump.git]
|
|
49
|
+
message_template = "Bump to {new_version}"
|
|
50
|
+
tag_template = "{new_version}"
|
|
51
|
+
|
|
52
|
+
[[tool.tbump.file]]
|
|
53
|
+
src = "pyproject.toml"
|
|
54
|
+
search = 'version = "{current_version}"'
|
|
55
|
+
|
|
56
|
+
[[tool.tbump.file]]
|
|
57
|
+
src = "src/rarg_numba_patterns/__init__.py"
|
|
58
|
+
search = '__version__ = "{current_version}"'
|
|
@@ -0,0 +1,47 @@
|
|
|
1
|
+
"""General-purpose numba intrinsics and literals patterns"""
|
|
2
|
+
|
|
3
|
+
from rarg_numba_patterns.intrinsics import (
|
|
4
|
+
accumulate_data,
|
|
5
|
+
atomic_rmw_intrinsic,
|
|
6
|
+
field_ptr,
|
|
7
|
+
item_ptr,
|
|
8
|
+
load_data,
|
|
9
|
+
overload_atomic_rmw,
|
|
10
|
+
overload_field_ptr,
|
|
11
|
+
overload_item_ptr,
|
|
12
|
+
)
|
|
13
|
+
from rarg_numba_patterns.literals import (
|
|
14
|
+
BooleanDatumLiteral,
|
|
15
|
+
Datum,
|
|
16
|
+
DatumLiteral,
|
|
17
|
+
FloatDatumLiteral,
|
|
18
|
+
IntegerDatumLiteral,
|
|
19
|
+
is_datum_literal,
|
|
20
|
+
LiteralStructRef,
|
|
21
|
+
Schema,
|
|
22
|
+
SchemaLiteral,
|
|
23
|
+
StringDatumLiteral,
|
|
24
|
+
)
|
|
25
|
+
|
|
26
|
+
__all__ = [
|
|
27
|
+
"accumulate_data",
|
|
28
|
+
"atomic_rmw_intrinsic",
|
|
29
|
+
"BooleanDatumLiteral",
|
|
30
|
+
"Datum",
|
|
31
|
+
"DatumLiteral",
|
|
32
|
+
"field_ptr",
|
|
33
|
+
"FloatDatumLiteral",
|
|
34
|
+
"IntegerDatumLiteral",
|
|
35
|
+
"is_datum_literal",
|
|
36
|
+
"item_ptr",
|
|
37
|
+
"LiteralStructRef",
|
|
38
|
+
"load_data",
|
|
39
|
+
"overload_atomic_rmw",
|
|
40
|
+
"overload_field_ptr",
|
|
41
|
+
"overload_item_ptr",
|
|
42
|
+
"Schema",
|
|
43
|
+
"SchemaLiteral",
|
|
44
|
+
"StringDatumLiteral",
|
|
45
|
+
]
|
|
46
|
+
|
|
47
|
+
__version__ = "0.0.1"
|
|
File without changes
|
|
@@ -0,0 +1,199 @@
|
|
|
1
|
+
"""This module contains code for some simple atomic spinlocks
|
|
2
|
+
|
|
3
|
+
See the discussion here
|
|
4
|
+
https://numba.discourse.group/t/phi-node-error-when-creating-a-while-loop-for-an-atomic-spinlock/3046
|
|
5
|
+
"""
|
|
6
|
+
|
|
7
|
+
import ctypes
|
|
8
|
+
import os
|
|
9
|
+
|
|
10
|
+
import numba
|
|
11
|
+
import numpy as np
|
|
12
|
+
from llvmlite import ir
|
|
13
|
+
from numba.core import cgutils, types
|
|
14
|
+
from numba.core.errors import RequireLiteralValue, TypingError
|
|
15
|
+
from numba.extending import intrinsic
|
|
16
|
+
|
|
17
|
+
if os.name == "posix":
|
|
18
|
+
YIELD_FUNCTION = "sched_yield"
|
|
19
|
+
elif os.name == "nt":
|
|
20
|
+
YIELD_FUNCTION = "SwitchToThread"
|
|
21
|
+
else:
|
|
22
|
+
raise NotImplementedError(f"atomic locks on OS {os.name}")
|
|
23
|
+
|
|
24
|
+
if not hasattr(ctypes.CDLL(None), YIELD_FUNCTION):
|
|
25
|
+
raise ImportError(
|
|
26
|
+
f"{YIELD_FUNCTION} not found in the default OS system libraries. "
|
|
27
|
+
f"atomic locks are not supported on this platform"
|
|
28
|
+
)
|
|
29
|
+
|
|
30
|
+
|
|
31
|
+
def _emit_lock_op(context, builder, operation, ptr, lock_type):
|
|
32
|
+
"""Emit IR for a lock or unlock operation on an already-resolved pointer."""
|
|
33
|
+
|
|
34
|
+
llvm_lock_type = context.get_value_type(lock_type.dtype)
|
|
35
|
+
|
|
36
|
+
ll_index_type = context.get_value_type(types.int64)
|
|
37
|
+
|
|
38
|
+
if operation.literal_value == "unlock":
|
|
39
|
+
builder.store_atomic(
|
|
40
|
+
ir.Constant(llvm_lock_type, 0),
|
|
41
|
+
ptr,
|
|
42
|
+
ordering="release",
|
|
43
|
+
align=context.get_abi_alignment(llvm_lock_type),
|
|
44
|
+
)
|
|
45
|
+
return ir.Constant(ir.IntType(1), 1)
|
|
46
|
+
|
|
47
|
+
# Spin until the lock is acquired
|
|
48
|
+
loop_cond = builder.append_basic_block(name="lock.while.cond")
|
|
49
|
+
loop_body = builder.append_basic_block(name="lock.while.body")
|
|
50
|
+
loop_end = builder.append_basic_block(name="lock.while.end")
|
|
51
|
+
|
|
52
|
+
# Get the OS yield function
|
|
53
|
+
try:
|
|
54
|
+
sched_yield = builder.module.globals[YIELD_FUNCTION]
|
|
55
|
+
except KeyError:
|
|
56
|
+
sched_yield_fnty = ir.FunctionType(ir.IntType(32), [])
|
|
57
|
+
sched_yield = ir.Function(builder.module, sched_yield_fnty, name=YIELD_FUNCTION)
|
|
58
|
+
|
|
59
|
+
start_block = builder.block
|
|
60
|
+
builder.branch(loop_cond)
|
|
61
|
+
|
|
62
|
+
with builder.goto_block(loop_cond):
|
|
63
|
+
count_phi = builder.phi(ll_index_type, name="lock.while.index")
|
|
64
|
+
xchng_result = builder.cmpxchg(
|
|
65
|
+
ptr,
|
|
66
|
+
ir.Constant(llvm_lock_type, 0),
|
|
67
|
+
ir.Constant(llvm_lock_type, 1),
|
|
68
|
+
ordering="acquire",
|
|
69
|
+
failordering="acquire",
|
|
70
|
+
)
|
|
71
|
+
success = builder.extract_value(xchng_result, 1)
|
|
72
|
+
pred = builder.icmp_signed("==", success, success.type(1))
|
|
73
|
+
builder.cbranch(pred, loop_end, loop_body)
|
|
74
|
+
|
|
75
|
+
with builder.goto_block(loop_body):
|
|
76
|
+
next_count = builder.add(count_phi, count_phi.type(1))
|
|
77
|
+
builder.call(sched_yield, [])
|
|
78
|
+
branch_block = builder.block
|
|
79
|
+
builder.branch(loop_cond)
|
|
80
|
+
|
|
81
|
+
count_phi.add_incoming(count_phi.type(0), start_block)
|
|
82
|
+
count_phi.add_incoming(next_count, branch_block)
|
|
83
|
+
|
|
84
|
+
builder.position_at_end(loop_end)
|
|
85
|
+
return ir.Constant(ir.IntType(1), 1)
|
|
86
|
+
|
|
87
|
+
|
|
88
|
+
@intrinsic(prefer_literal=True)
|
|
89
|
+
def atomic_lock(typingctx, lock_type):
|
|
90
|
+
"""Allocates an integer suitable for use as an atomic lock
|
|
91
|
+
on the stack and returns a pointer to that integer"""
|
|
92
|
+
|
|
93
|
+
if not isinstance(lock_type, (types.NumberClass, types.DType)):
|
|
94
|
+
raise TypingError(f"{lock_type} is not a NumberClass or DType")
|
|
95
|
+
|
|
96
|
+
def codegen(context, builder, signature, args):
|
|
97
|
+
llvm_lock_type = context.get_value_type(lock_type.dtype)
|
|
98
|
+
ptr = cgutils.alloca_once(builder, llvm_lock_type)
|
|
99
|
+
builder.store(ir.Constant(llvm_lock_type, 0), ptr)
|
|
100
|
+
return ptr
|
|
101
|
+
|
|
102
|
+
return types.CPointer(lock_type.dtype)(lock_type), codegen
|
|
103
|
+
|
|
104
|
+
|
|
105
|
+
@intrinsic(prefer_literal=True)
|
|
106
|
+
def lock_int_op(typingctx, lock: types.CPointer, operation: types.StringLiteral):
|
|
107
|
+
"""Performs an atomic lock/lock on integer located by a pointer"""
|
|
108
|
+
if not isinstance(operation, types.StringLiteral) or operation.literal_value not in {
|
|
109
|
+
"lock",
|
|
110
|
+
"unlock",
|
|
111
|
+
}:
|
|
112
|
+
raise RequireLiteralValue(
|
|
113
|
+
f"'operation' {operation} must be a StringLiteral "
|
|
114
|
+
f"set to either lock or unlock"
|
|
115
|
+
)
|
|
116
|
+
|
|
117
|
+
if not (isinstance(lock, types.CPointer) and isinstance(lock.dtype, types.Integer)):
|
|
118
|
+
raise TypingError(f"lock {lock} must be an CPointer to an integer")
|
|
119
|
+
|
|
120
|
+
sig = types.bool(lock, operation)
|
|
121
|
+
|
|
122
|
+
def codegen(context, builder, signature, args):
|
|
123
|
+
lock, _ = args
|
|
124
|
+
lock_type, _ = signature.args
|
|
125
|
+
return _emit_lock_op(context, builder, operation, lock, lock_type)
|
|
126
|
+
|
|
127
|
+
return sig, codegen
|
|
128
|
+
|
|
129
|
+
|
|
130
|
+
@intrinsic(prefer_literal=True)
|
|
131
|
+
def lock_array_op(
|
|
132
|
+
typingctx, lock: types.Array, idx: types.UniTuple, operation: types.StringLiteral
|
|
133
|
+
):
|
|
134
|
+
"""Performs an atomic lock/unlock operation at the given index in an integer array"""
|
|
135
|
+
if not isinstance(operation, types.StringLiteral) or operation.literal_value not in {
|
|
136
|
+
"lock",
|
|
137
|
+
"unlock",
|
|
138
|
+
}:
|
|
139
|
+
raise RequireLiteralValue(
|
|
140
|
+
f"'operation' {operation} must be a StringLiteral "
|
|
141
|
+
f"set to either lock or unlock"
|
|
142
|
+
)
|
|
143
|
+
|
|
144
|
+
if not isinstance(lock, types.Array) or not isinstance(lock.dtype, types.Integer):
|
|
145
|
+
raise TypingError(f"lock {lock} must be an Array of integers")
|
|
146
|
+
|
|
147
|
+
if (
|
|
148
|
+
not isinstance(idx, types.UniTuple)
|
|
149
|
+
or not isinstance(idx.dtype, types.Integer)
|
|
150
|
+
or len(idx) != lock.ndim
|
|
151
|
+
):
|
|
152
|
+
raise TypingError(f"idx {idx} must be a Tuple of length {lock.ndim} integers")
|
|
153
|
+
|
|
154
|
+
sig = types.bool(lock, idx, operation)
|
|
155
|
+
|
|
156
|
+
def codegen(context, builder, signature, args):
|
|
157
|
+
lock, idx, _ = args
|
|
158
|
+
lock_type, idx_type, _ = signature.args
|
|
159
|
+
lock_array = context.make_array(lock_type)(context, builder, lock)
|
|
160
|
+
native_idx = [builder.extract_value(idx, i) for i in range(len(idx_type))]
|
|
161
|
+
out_ptr = builder.gep(lock_array.data, native_idx)
|
|
162
|
+
return _emit_lock_op(context, builder, operation, out_ptr, lock_type)
|
|
163
|
+
|
|
164
|
+
return sig, codegen
|
|
165
|
+
|
|
166
|
+
|
|
167
|
+
if __name__ == "__main__":
|
|
168
|
+
"""Test script"""
|
|
169
|
+
|
|
170
|
+
@numba.njit(nogil=True)
|
|
171
|
+
def lock_index():
|
|
172
|
+
ll = numba.typed.List()
|
|
173
|
+
ll.append(atomic_lock(np.dtype(np.uint8)))
|
|
174
|
+
ll.append(atomic_lock(np.dtype(np.uint8)))
|
|
175
|
+
ll.append(atomic_lock(np.dtype(np.uint8)))
|
|
176
|
+
lock_int_op(ll[0], "lock")
|
|
177
|
+
lock_int_op(ll[0], "unlock")
|
|
178
|
+
|
|
179
|
+
return ll
|
|
180
|
+
|
|
181
|
+
print(len(lock_index()))
|
|
182
|
+
|
|
183
|
+
if True:
|
|
184
|
+
|
|
185
|
+
@numba.njit(nogil=True)
|
|
186
|
+
def lock_index(a, i):
|
|
187
|
+
return lock_array_op(a, i, "lock")
|
|
188
|
+
|
|
189
|
+
@numba.njit(nogil=True)
|
|
190
|
+
def unlock_index(a, i):
|
|
191
|
+
return lock_array_op(a, i, "unlock")
|
|
192
|
+
|
|
193
|
+
locks = np.full(10, 0, np.int32)
|
|
194
|
+
|
|
195
|
+
print(lock_index(locks, (0,)))
|
|
196
|
+
print(lock_index(locks, (1,)))
|
|
197
|
+
print(unlock_index(locks, (0,)))
|
|
198
|
+
print(lock_index(locks, (0,)))
|
|
199
|
+
print(unlock_index(locks, (1,)))
|