diff --git a/.github/workflows/examples_kokkos.yml b/.github/workflows/examples_kokkos.yml new file mode 100644 index 00000000..09c21359 --- /dev/null +++ b/.github/workflows/examples_kokkos.yml @@ -0,0 +1,48 @@ +name: examples/kokkos + +on: + push: + branches: + - main + pull_request: + branches: + - main + +jobs: + test_pykokkos: + strategy: + matrix: + platform: [ubuntu-latest] + python-version: ["3.13"] + runs-on: ${{ matrix.platform }} + steps: + - uses: actions/checkout@v3 + - name: Set up Python ${{ matrix.python-version }} + uses: actions/setup-python@v4 + with: + python-version: ${{ matrix.python-version }} + - name: Install dependencies + run: | + python -m pip install --upgrade pip + python -m pip install --upgrade numpy mypy==1.0.1 cmake pytest pybind11 scikit-build patchelf + - name: Install pykokkos-base + run: | + python install_base.py install -- -DENABLE_LAYOUTS=ON -DENABLE_MEMORY_TRAITS=OFF -DENABLE_VIEW_RANKS=2 + - name: Install pykokkos + run: | + python -m pip install . + - name: mypy check + run: | + mypy pykokkos + - name: Run all kokkos examples + run: | + for f in examples/kokkos/*.py; do + if [[ "$f" == *_cuda.py ]]; then + echo "Ignoring cuda scripts $f" + elif [[ "$f" == *_lambda.py ]]; then + echo "Ignoring lambda script $f, see https://github.com/kokkos/pykokkos/issues/346" + else + echo "Running $f..." + python "$f" || exit 1 + fi + done diff --git a/README.md b/README.md index 6f0ee18f..8d160f55 100644 --- a/README.md +++ b/README.md @@ -4,6 +4,7 @@ [![Documentation](https://github.com/kokkos/pykokkos/actions/workflows/documentation.yml/badge.svg)](https://github.com/kokkos/pykokkos/actions/workflows/documentation.yml) [![Linux](https://github.com/kokkos/pykokkos/actions/workflows/build_linux.yml/badge.svg)](https://github.com/kokkos/pykokkos/actions/workflows/build_linux.yml) [![MacOS](https://github.com/kokkos/pykokkos/actions/workflows/build_macos.yml/badge.svg)](https://github.com/kokkos/pykokkos/actions/workflows/build_macos.yml) +[![examples/kokkos](https://github.com/kokkos/pykokkos/actions/workflows/examples_kokkos.yml/badge.svg)](https://github.com/kokkos/pykokkos/actions/workflows/examples_kokkos.yml) PyKokkos is a framework for writing high-performance Python code similar to Numba. In contrast to Numba, PyKokkos kernels are primarily diff --git a/examples/kokkos/01_hello_world.py b/examples/kokkos/01_hello_world.py index 7c9ee228..3764663f 100644 --- a/examples/kokkos/01_hello_world.py +++ b/examples/kokkos/01_hello_world.py @@ -8,7 +8,6 @@ def hello(i: int): def main(): N: int = 10 - pk.set_default_space(pk.ExecutionSpace.OpenMP) pk.parallel_for(N, hello) diff --git a/examples/kokkos/01_hello_world_lambda.py b/examples/kokkos/01_hello_world_lambda.py index 05d5104a..8a2a4208 100644 --- a/examples/kokkos/01_hello_world_lambda.py +++ b/examples/kokkos/01_hello_world_lambda.py @@ -1,15 +1,10 @@ import pykokkos as pk -@pk.workload -class HelloWorld: - def __init__(self, n): - self.N: int = n - - @pk.main - def run(self): - pk.parallel_for(self.N, lambda i: pk.printf("Hello from i = %i\n", i)) +def main(): + N: int = 10 + pk.parallel_for(N, lambda i: pk.printf("Hello from i = %d\n", i)) if __name__ == "__main__": - pk.execute(pk.ExecutionSpace.OpenMP, HelloWorld(10)) + main() diff --git a/examples/kokkos/02_simple_reduce.py b/examples/kokkos/02_simple_reduce.py index cf14bccc..f063345c 100644 --- a/examples/kokkos/02_simple_reduce.py +++ b/examples/kokkos/02_simple_reduce.py @@ -9,7 +9,6 @@ def squaresum(i: int, acc, values): def main(): N: int = 10 - pk.set_default_space(pk.ExecutionSpace.OpenMP) # Create array with squares values = np.array([i * i for i in range(N)], dtype=np.int32) diff --git a/examples/kokkos/02_simple_reduce_lambda.py b/examples/kokkos/02_simple_reduce_lambda.py index b54c92ea..fbc35142 100644 --- a/examples/kokkos/02_simple_reduce_lambda.py +++ b/examples/kokkos/02_simple_reduce_lambda.py @@ -1,20 +1,13 @@ import pykokkos as pk -@pk.workload -class SquareSum: - def __init__(self, n): - self.N: int = n - self.total: int = 0 +def main(): + N: int = 10 - @pk.main - def run(self): - self.total = pk.parallel_reduce(self.N, lambda i, acc: acc + i * i) + total = pk.parallel_reduce(N, lambda i, acc: acc + i * i) - @pk.callback - def results(self): - print("Sum:", self.total) + print("Sum:", total) if __name__ == "__main__": - pk.execute(pk.ExecutionSpace.OpenMP, SquareSum(10)) + main() diff --git a/examples/kokkos/03_simple_view.py b/examples/kokkos/03_simple_view.py index 5e9aa22a..0efc817b 100644 --- a/examples/kokkos/03_simple_view.py +++ b/examples/kokkos/03_simple_view.py @@ -14,7 +14,6 @@ def my_reduction(i: int, accumulator: pk.Acc[pk.double], a: pk.View2D[pk.int32]) def main(): N: int = 10 - pk.set_default_space(pk.ExecutionSpace.OpenMP) a: pk.View2D[pk.int32] = pk.View([N, 3], pk.int32) diff --git a/examples/kokkos/03_simple_view_lambda.py b/examples/kokkos/03_simple_view_lambda.py index ea9dbe8e..a6428d50 100644 --- a/examples/kokkos/03_simple_view_lambda.py +++ b/examples/kokkos/03_simple_view_lambda.py @@ -1,33 +1,23 @@ +import numpy as np import pykokkos as pk -@pk.workload -class SimpleView: - def __init__(self, n): - self.N: int = n - self.total: int = 0 - self.a: pk.View2D[pk.int32] = pk.View([self.N, 3], pk.int32) +def main(): + N = 10 - @pk.callback - def results(self): - for row in self.a: - print(row) - print("\nResult is", self.total) + i = np.arange(1, N + 1, dtype=np.int32) + j = np.arange(1, 3 + 1, dtype=np.int32) + a = i.reshape(-1, 1) ** j.reshape(1, -1) + print(a) - @pk.main - def run(self): - pk.parallel_for(self.N, self.initialize_view) - self.total = pk.parallel_reduce( - self.N, - lambda i, accumulator: accumulator - + self.a[i][0] * self.a[i][1] / (self.a[i][2]), - ) + total: int = pk.parallel_reduce( + N, lambda u, acc: acc + a[i][0] * a[i][1] / a[i][2], a=a + ) - @pk.workunit - def initialize_view(self, i: int): - for j in range(3): - self.a[i][j] = (i + 1) ** (j + 1) + for row in a: + print(row) + print("\nResult is", total) if __name__ == "__main__": - pk.execute(pk.ExecutionSpace.OpenMP, SimpleView(10)) + main() diff --git a/examples/kokkos/04_simple_memoryspaces.py b/examples/kokkos/04_simple_memoryspaces.py index 7ffb61f9..893116ac 100644 --- a/examples/kokkos/04_simple_memoryspaces.py +++ b/examples/kokkos/04_simple_memoryspaces.py @@ -8,7 +8,6 @@ def reduction(i: int, acc: pk.Acc[pk.double], a: pk.View2D[pk.int32]): def main(): N: int = 10 - pk.set_default_space(pk.ExecutionSpace.OpenMP) a: pk.View2D[pk.int32] = pk.View([N, 3], pk.int32) diff --git a/examples/kokkos/04_simple_memoryspaces_lambda.py b/examples/kokkos/04_simple_memoryspaces_lambda.py index b301350e..ac920d2d 100644 --- a/examples/kokkos/04_simple_memoryspaces_lambda.py +++ b/examples/kokkos/04_simple_memoryspaces_lambda.py @@ -1,30 +1,21 @@ +import numpy as np import pykokkos as pk -@pk.workload -class SimpleSpaces: - def __init__(self, n): - self.N: int = n - self.sum: int = 0 - self.a: pk.View2D[pk.int32] = pk.View([n, 3], pk.int32) - for i in range(n): - for j in range(3): - self.a[i][j] = i * n + j +def main(): + N = 10 - @pk.main - def run(self): - self.sum = pk.parallel_reduce( - self.N, - lambda i, accumulator: accumulator - + self.a[i][0] - - self.a[i][1] - + self.a[i][2], - ) + # Initialize the array + i = np.arange(N, dtype=np.int32) + j = np.arange(3, dtype=np.int32) + a = i.reshape(-1, 1) * N + j.reshape(1, -1) - @pk.callback - def use_results(self): - print(self.sum) + sum_result = pk.parallel_reduce( + N, lambda i, acc: acc + a[i][0] - a[i][1] + a[i][2], a=a + ) + + print(sum_result) if __name__ == "__main__": - pk.execute(pk.ExecutionSpace.OpenMP, SimpleSpaces(10)) + main() diff --git a/examples/kokkos/05_simple_atomics.py b/examples/kokkos/05_simple_atomics.py index 9bb833b1..58f381f3 100644 --- a/examples/kokkos/05_simple_atomics.py +++ b/examples/kokkos/05_simple_atomics.py @@ -29,7 +29,6 @@ def findprimes( def simple_atomics(): N: int = 100 - pk.set_default_space(pk.ExecutionSpace.OpenMP) data: pk.View1D[pk.int32] = pk.View([N], pk.int32) result: pk.View1D[pk.int32] = pk.View([N], pk.int32) diff --git a/examples/kokkos/README.md b/examples/kokkos/README.md index 87545332..c1e479cc 100644 --- a/examples/kokkos/README.md +++ b/examples/kokkos/README.md @@ -1,2 +1,4 @@ This directory contains examples translated from the main Kokkos repository: https://github.com/kokkos/kokkos/tree/develop/example/tutorial + +Except for `inclusive_scan_team_cuda.py`, each example executes on the default Kokkos execution space specified in the Kokkos build diff --git a/examples/kokkos/add1.py b/examples/kokkos/add1.py index 78dbb158..71c92683 100644 --- a/examples/kokkos/add1.py +++ b/examples/kokkos/add1.py @@ -9,7 +9,6 @@ def add1(i: int, a: pk.View1D[pk.int32]): def main(): n: int = 100 * 1000 N: int = n - pk.set_default_space(pk.ExecutionSpace.OpenMP) a: pk.View1D[pk.int32] = pk.View([N], pk.int32) diff --git a/examples/kokkos/add1_lambda.py b/examples/kokkos/add1_lambda.py index f422fa0e..bdca1d79 100644 --- a/examples/kokkos/add1_lambda.py +++ b/examples/kokkos/add1_lambda.py @@ -1,26 +1,24 @@ +import numpy as np import pykokkos as pk -@pk.workload -class AddOne: - def __init__(self, n): - self.N: int = n - self.a: pk.View1D[pk.int32] = pk.View([n], pk.int32) +@pk.workunit +def add1(i: int, a: pk.View1D[pk.int32]): + a[i] += 1 - for i in range(self.N): - self.a[i] = 2 - print(f"Initialized view: [{self.a[0]}, ... repeats {n-1} times]") - @pk.main - def run(self): - y: int = 1 - pk.parallel_for(self.N, lambda i: self.a[i] + y, self.a) +def main(): + n = 100 * 1000 + N = n + + # Initialize the array + a = 2 * np.ones(N, dtype=np.int32) + print(f"Initialized view: [{a[0]}, ... repeats {n-1} times]") - @pk.callback - def results(self): - print(f"Results: [{self.a[0]}, ... repeats {n-1} times]") + pk.parallel_for(N, lambda i: a[i] + 1, a=a) + + print(f"Results: [{a[0]}, ... repeats {n-1} times]") if __name__ == "__main__": - n = 100 * 1000 - pk.execute(pk.ExecutionSpace.OpenMP, AddOne(n)) + main() diff --git a/examples/kokkos/inclusive_scan_team.py b/examples/kokkos/inclusive_scan_team.py index 016d19a5..435e1d78 100644 --- a/examples/kokkos/inclusive_scan_team.py +++ b/examples/kokkos/inclusive_scan_team.py @@ -34,12 +34,12 @@ def main(): num_teams = (N + team_size - 1) // team_size view = np.zeros([N], dtype=np.int32) - p_init = pk.RangePolicy(pk.ExecutionSpace.OpenMP, 0, N) + p_init = pk.RangePolicy(0, N) pk.parallel_for(p_init, init_data, view=view) print(f"Total elements: {N}, Team size: {team_size}, Number of teams: {num_teams}") - team_policy = pk.TeamPolicy(pk.ExecutionSpace.OpenMP, num_teams, team_size) + team_policy = pk.TeamPolicy(num_teams, team_size) print("Running kernel...") pk.parallel_for(team_policy, team_scan, view=view, team_size=team_size) diff --git a/examples/kokkos/math_functions.py b/examples/kokkos/math_functions.py index 92118a6a..7bc2dc59 100644 --- a/examples/kokkos/math_functions.py +++ b/examples/kokkos/math_functions.py @@ -11,7 +11,6 @@ def my_calculation(i: int, a: pk.View1D[pk.int32], N: int): def main(): n: int = 10 N: int = n - pk.set_default_space(pk.ExecutionSpace.OpenMP) a: pk.View1D[pk.int32] = pk.View([N], pk.int32) diff --git a/examples/kokkos/matrix_sum.py b/examples/kokkos/matrix_sum.py index 61a7362b..4df0b45a 100644 --- a/examples/kokkos/matrix_sum.py +++ b/examples/kokkos/matrix_sum.py @@ -15,7 +15,6 @@ def final_sum(i: int, accumulator: pk.Acc[pk.double], mat: pk.View2D[pk.int32]): def main(): r: int = 5 c: int = 10 - pk.set_default_space(pk.ExecutionSpace.OpenMP) mat: pk.View2D[pk.int32] = pk.View([r, c], pk.int32) diff --git a/examples/kokkos/multi_scratch_team.py b/examples/kokkos/multi_scratch_team_cuda.py similarity index 100% rename from examples/kokkos/multi_scratch_team.py rename to examples/kokkos/multi_scratch_team_cuda.py diff --git a/examples/kokkos/random_sum.py b/examples/kokkos/random_sum.py index 893860e3..86ee1ead 100644 --- a/examples/kokkos/random_sum.py +++ b/examples/kokkos/random_sum.py @@ -10,7 +10,6 @@ def my_reduction(i: int, accumulator: pk.Acc[pk.int32], a: pk.View1D[pk.int32]): def main(): n: int = 10 N: int = n - pk.set_default_space(pk.ExecutionSpace.OpenMP) a: pk.View1D[pk.int32] = pk.View([N], pk.int32) diff --git a/examples/kokkos/scan_standalone.py b/examples/kokkos/scan.py similarity index 90% rename from examples/kokkos/scan_standalone.py rename to examples/kokkos/scan.py index c19ddf08..ba957d19 100644 --- a/examples/kokkos/scan_standalone.py +++ b/examples/kokkos/scan.py @@ -18,7 +18,7 @@ def run() -> None: N = 10 A = np.zeros([N], dtype=np.int32) - p = pk.RangePolicy(pk.ExecutionSpace.OpenMP, 0, N) + p = pk.RangePolicy(0, N) pk.parallel_for(p, init, view=A) timer = pk.Timer() diff --git a/examples/kokkos/scan_functor.py b/examples/kokkos/scan_functor.py deleted file mode 100644 index f85e09ba..00000000 --- a/examples/kokkos/scan_functor.py +++ /dev/null @@ -1,34 +0,0 @@ -import pykokkos as pk - - -@pk.workunit -def init(i: int, A: pk.View1D[pk.int32]): - A[i] = i - - -@pk.workunit -def scan(i: int, acc: pk.Acc[pk.double], last_pass: bool, A: pk.View1D[pk.int32]): - acc += A[i] - if last_pass: - A[i] = acc - - -def run() -> None: - N: int = 10 - pk.set_default_space(pk.ExecutionSpace.OpenMP) - - A: pk.View1D[pk.int32] = pk.View([N], pk.int32) - - p = pk.RangePolicy(pk.ExecutionSpace.OpenMP, 0, N) - - pk.parallel_for(p, init, A=A) - - timer = pk.Timer() - result = pk.parallel_scan(p, scan, A=A) - timer_result = timer.seconds() - - print(f"{A} total={result} time({timer_result})") - - -if __name__ == "__main__": - run() diff --git a/examples/kokkos/scan_workload.py b/examples/kokkos/scan_workload.py deleted file mode 100644 index e51235ad..00000000 --- a/examples/kokkos/scan_workload.py +++ /dev/null @@ -1,32 +0,0 @@ -import pykokkos as pk - - -@pk.workunit -def init(i: int, A: pk.View1D[pk.int32]): - A[i] = i - - -@pk.workunit -def scan(i: int, acc: pk.Acc[pk.double], last_pass: bool, A: pk.View1D[pk.int32]): - acc += A[i] - if last_pass: - A[i] = acc - - -def run() -> None: - N: int = 10 - pk.set_default_space(pk.ExecutionSpace.OpenMP) - - A: pk.View1D[pk.int32] = pk.View([N], pk.int32) - - pk.parallel_for(N, init, A=A) - - timer = pk.Timer() - result: int = pk.parallel_scan(N, scan, A=A) - timer_result: float = timer.seconds() - - print(f"{A} total={result} time({timer_result})") - - -if __name__ == "__main__": - run() diff --git a/pykokkos/interface/parallel_dispatch.py b/pykokkos/interface/parallel_dispatch.py index ee8d4394..ddc01684 100644 --- a/pykokkos/interface/parallel_dispatch.py +++ b/pykokkos/interface/parallel_dispatch.py @@ -169,6 +169,19 @@ def handle_args(is_for: bool, *args) -> HandledArgs: f"ERROR: workunit expected to be type 'Callable' or 'List[Callable]', got '{workunit}' of type '{type(workunit)}'" ) + # check that workunit is not a lambda + if isinstance(workunit, list): + for wu in workunit: + if wu.__name__ == "": + raise TypeError( + f"ERROR: lambda functions are not currently supported for translation, see issue https://github.com/kokkos/pykokkos/issues/346" + ) + else: + if workunit.__name__ == "": + raise TypeError( + f"ERROR: lambda functions are not currently supported for translation, see issue https://github.com/kokkos/pykokkos/issues/346" + ) + return HandledArgs(name, policy, workunit, view, initial_value)