Skip to content

Commit effcd59

Browse files
nayyirahsannayyirahsan
committed
Merge branch 'nayyirahsan/ufunc-removal' of https://github.com/nayyirahsan/pykokkos into nayyirahsan/ufunc-removal
2 parents 4159771 + 8fb75d7 commit effcd59

29 files changed

+500
-464
lines changed

‎.github/workflows/build_linux.yml‎

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,39 @@
1+
name: Linux
2+
3+
on:
4+
push:
5+
branches:
6+
- main
7+
pull_request:
8+
branches:
9+
- main
10+
11+
jobs:
12+
test_pykokkos:
13+
strategy:
14+
matrix:
15+
platform: [ubuntu-latest, ubuntu-24.04-arm]
16+
python-version: ["3.13"]
17+
runs-on: ${{ matrix.platform }}
18+
steps:
19+
- uses: actions/checkout@v3
20+
- name: Set up Python ${{ matrix.python-version }}
21+
uses: actions/setup-python@v4
22+
with:
23+
python-version: ${{ matrix.python-version }}
24+
- name: Install dependencies
25+
run: |
26+
python -m pip install --upgrade pip
27+
python -m pip install --upgrade numpy mypy==1.0.1 cmake pytest pybind11 scikit-build patchelf
28+
- name: Install pykokkos-base
29+
run: |
30+
python install_base.py install -- -DENABLE_LAYOUTS=ON -DENABLE_MEMORY_TRAITS=OFF -DENABLE_VIEW_RANKS=1
31+
- name: Install pykokkos
32+
run: |
33+
python -m pip install .
34+
- name: mypy check
35+
run: |
36+
mypy pykokkos
37+
- name: import pykokkos
38+
run: |
39+
python examples/kokkos-tutorials/standalone/01.py

‎.github/workflows/build_macos.yml‎

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,39 @@
1+
name: MacOS
2+
3+
on:
4+
push:
5+
branches:
6+
- main
7+
pull_request:
8+
branches:
9+
- main
10+
11+
jobs:
12+
test_pykokkos:
13+
strategy:
14+
matrix:
15+
platform: [macos-latest]
16+
python-version: ["3.13"]
17+
runs-on: ${{ matrix.platform }}
18+
steps:
19+
- uses: actions/checkout@v3
20+
- name: Set up Python ${{ matrix.python-version }}
21+
uses: actions/setup-python@v4
22+
with:
23+
python-version: ${{ matrix.python-version }}
24+
- name: Install dependencies
25+
run: |
26+
python -m pip install --upgrade pip
27+
python -m pip install --upgrade numpy mypy==1.0.1 cmake pytest pybind11 scikit-build patchelf
28+
- name: Install pykokkos-base
29+
run: |
30+
python install_base.py install -- -DENABLE_LAYOUTS=ON -DENABLE_MEMORY_TRAITS=OFF -DENABLE_VIEW_RANKS=1
31+
- name: Install pykokkos
32+
run: |
33+
python -m pip install .
34+
- name: mypy check
35+
run: |
36+
mypy pykokkos
37+
- name: import pykokkos
38+
run: |
39+
python examples/kokkos-tutorials/standalone/01.py

‎README.md‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,8 @@
22
# PyKokkos
33
[![Python Testing](https://github.com/kokkos/pykokkos/actions/workflows/main_ci.yml/badge.svg)](https://github.com/kokkos/pykokkos/actions/workflows/main_ci.yml)
44
[![Documentation](https://github.com/kokkos/pykokkos/actions/workflows/documentation.yml/badge.svg)](https://github.com/kokkos/pykokkos/actions/workflows/documentation.yml)
5+
[![Linux](https://github.com/kokkos/pykokkos/actions/workflows/build_linux.yml/badge.svg)](https://github.com/kokkos/pykokkos/actions/workflows/build_linux.yml)
6+
[![MacOS](https://github.com/kokkos/pykokkos/actions/workflows/build_macos.yml/badge.svg)](https://github.com/kokkos/pykokkos/actions/workflows/build_macos.yml)
57

68
PyKokkos is a framework for writing high-performance Python code
79
similar to Numba. In contrast to Numba, PyKokkos kernels are primarily

‎base/CMakeLists.txt‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -37,8 +37,8 @@ if(NOT TARGET pybind11::module)
3737
MESSAGE(STATUS "Fetching pybind11")
3838
FETCHCONTENT_DECLARE(
3939
pybind11
40-
URL https://github.com/pybind/pybind11/archive/refs/heads/v3.0.zip
41-
URL_HASH SHA256=1816ea10bf7de362fb7864b131cc90bed6fc393feaf7e27b0ed5a03987bedd10
40+
URL https://github.com/pybind/pybind11/archive/refs/tags/v3.0.0.zip
41+
URL_HASH SHA256=dfe152af2f454a9d8cd771206c014aecb8c3977822b5756123f29fd488648334
4242
)
4343
FETCHCONTENT_MAKEAVAILABLE(pybind11)
4444
endif()
@@ -159,7 +159,7 @@ SET_TARGET_PROPERTIES(libpykokkos PROPERTIES
159159
# configure the rpath: <RELATIVE>:<CWD>:<FULL>
160160
IF(APPLE)
161161
SET_TARGET_PROPERTIES(libpykokkos PROPERTIES
162-
MACOSX_RPATH "@loader_path/${LIB_RELPATH}:@loader_path:${_Kokkos_LIBDIR}:${CMAKE_INSTALL_RPATH}")
162+
INSTALL_RPATH "@loader_path/${LIB_RELPATH};@loader_path;${_Kokkos_LIBDIR};${CMAKE_INSTALL_RPATH}")
163163
ELSEIF(UNIX)
164164
SET_TARGET_PROPERTIES(libpykokkos PROPERTIES
165165
INSTALL_RPATH "\$ORIGIN/${LIB_RELPATH}:\$ORIGIN:${_Kokkos_LIBDIR}:${CMAKE_INSTALL_RPATH}")

‎base/environment.yml‎

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,4 @@ dependencies:
55
- cmake>=3.16
66
- scikit-build
77
- numpy
8-
- gcc=11.4.0
9-
- gxx=11.4.0
10-
8+

‎docs/environment_variables.rst‎

Lines changed: 48 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,48 @@
1+
2+
Environment Variables
3+
=====================
4+
5+
PyKokkos behavior can be controlled through the following environment variables:
6+
7+
.. list-table::
8+
:header-rows: 1
9+
:widths: 25 20 55
10+
11+
* - Variable
12+
- Value
13+
- Description
14+
* - ``PK_FORMAT``
15+
- any (presence triggers)
16+
- Runs ``clang-format`` on intermediate C++ sources.
17+
* - ``PK_FUSION``
18+
- ``"naive"``
19+
- Enables automatic kernel fusion. See :doc:`fusion`.
20+
* - ``PK_KOKKOS_LIB_PATH``
21+
- path to directory
22+
- Overrides search path for the compiled pykokkos-base
23+
``lib/`` or ``lib64/`` directory.
24+
* - ``PK_KOKKOS_INTERFACE``
25+
- Kokkos version string
26+
- Selects a specific version of the Kokkos interface to use when
27+
multiple versions are available.
28+
29+
Usage Examples
30+
--------------
31+
32+
Enable C++ code formatting for generated kernels:
33+
34+
.. code-block:: bash
35+
36+
export PK_FORMAT=1
37+
38+
Enable naive kernel fusion:
39+
40+
.. code-block:: bash
41+
42+
export PK_FUSION="naive"
43+
44+
Specify a custom Kokkos library path:
45+
46+
.. code-block:: bash
47+
48+
export PK_KOKKOS_LIB_PATH="/path/to/pykokkos-base/lib"

‎docs/index.rst‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -53,6 +53,7 @@ interest to those already familiar with the Kokkos framework.
5353
workunits
5454
ndarrays
5555
fusion
56+
environment_variables
5657
examples
5758

5859
..

‎examples/kokkos-tutorials/workload/team_scratch_memory.py‎

Lines changed: 0 additions & 129 deletions
This file was deleted.

‎examples/kokkos-tutorials/workload/team_scratch_workunit.py‎

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@
88
"""
99

1010

11-
@pk.workunit
11+
@pk.workunit(scratch=[(float, lambda p: p.M)])
1212
def yAx(team_member, acc: pk.Acc[float], y, x, A, M, N):
1313
e: int = team_member.league_rank()
1414
s_x: pk.ScratchView1D[float] = pk.ScratchView1D(team_member.team_scratch(0), M)
@@ -48,7 +48,7 @@ def single_closure():
4848
nrepeat: int = values[4]
4949
fill: bool = values[-1]
5050

51-
space = pk.ExecutionSpace.OpenMP
51+
space = pk.ExecutionSpace.Cuda
5252
pk.set_default_space(space)
5353

5454
# Initialize data
@@ -71,13 +71,12 @@ def single_closure():
7171
acc = 0
7272

7373
timer = pk.Timer()
74-
scratch_size: int = pk.ScratchView1D[float].shmem_size(M)
7574
print(f"Before: {N} | {M} | {E}")
7675

7776
for i in range(nrepeat):
7877
result = pk.parallel_reduce(
7978
"team_scratch_workunit",
80-
pk.TeamPolicy(E, "auto", 32).set_scratch_size(0, pk.PerTeam(scratch_size)),
79+
pk.TeamPolicy(E, "auto", 32),
8180
yAx,
8281
acc=acc,
8382
y=y,

‎examples/kokkos/inclusive_scan_team.py‎

Lines changed: 4 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -8,9 +8,8 @@ def init_data(i: int, view):
88

99

1010
# Test inclusive_scan with scratch memory
11-
@pk.workunit
12-
def team_scan(team_member: pk.TeamMember, view):
13-
team_size: int = team_member.team_size()
11+
@pk.workunit(scratch=[(int, lambda p: p.team_size)])
12+
def team_scan(team_member: pk.TeamMember, view, team_size):
1413
offset: int = team_member.league_rank() * team_size
1514
localIdx: int = team_member.team_rank()
1615
globalIdx: int = offset + localIdx
@@ -31,7 +30,7 @@ def team_scan(team_member: pk.TeamMember, view):
3130

3231
def main():
3332
N = 64
34-
team_size = 32
33+
team_size = 8
3534
num_teams = (N + team_size - 1) // team_size
3635

3736
view = np.zeros([N], dtype=np.int32)
@@ -42,12 +41,8 @@ def main():
4241

4342
team_policy = pk.TeamPolicy(pk.ExecutionSpace.OpenMP, num_teams, team_size)
4443

45-
# Set scratch size for current team policy
46-
scratch_size = pk.ScratchView1D[int].shmem_size(team_size)
47-
team_policy.set_scratch_size(0, pk.PerTeam(scratch_size))
48-
4944
print("Running kernel...")
50-
pk.parallel_for(team_policy, team_scan, view=view)
45+
pk.parallel_for(team_policy, team_scan, view=view, team_size=team_size)
5146
print(f"View, splitted by two groups of size = {team_size}")
5247
print(view)
5348

0 commit comments

Comments
 (0)