make examples more helpful

Author: nayyirahsan

examples/graphs/bfs_bottomup.py

@@ -1,11 +1,65 @@
+import argparse
+from collections import deque
+
+import numpy as np_cpu
+
 import pykokkos as pk
 
 if pk.get_default_space() in pk.DeviceExecutionSpace:
     import cupy as np
 else:
     import numpy as np
 
-import argparse
+
+def _view_to_numpy_host(x):
+    """Host NumPy array for comparisons (handles Cupy-backed views)."""
+    if hasattr(x, "get"):
+        return np_cpu.asarray(x.get())
+    return np_cpu.asarray(x)
+
+
+def reference_grid_bfs_distances(N: int, M: int, mat) -> np_cpu.ndarray:
+    """
+    Shortest hop count (4-neighbor grid) from every cell to any cell with mat==0.
+    Same graph as the PyKokkos workunits: vertices are grid cells, edges to N/E/S/W
+    neighbors. Multi-source BFS using a queue, following the standard pattern in
+    https://www.geeksforgeeks.org/python/python-program-for-breadth-first-search-or-bfs-for-a-graph/
+    """
+    mat_h = _view_to_numpy_host(mat)
+    dist = np_cpu.full(N * M, -1, dtype=np_cpu.int32)
+    q = deque()
+    for r in range(N):
+        for c in range(M):
+            if mat_h[r, c] == 0:
+                idx = r * M + c
+                dist[idx] = 0
+                q.append((r, c))
+    while q:
+        r, c = q.popleft()
+        d = int(dist[r * M + c])
+        for dr, dc in ((-1, 0), (1, 0), (0, -1), (0, 1)):
+            nr, nc = r + dr, c + dc
+            if 0 <= nr < N and 0 <= nc < M:
+                ni = nr * M + nc
+                if dist[ni] == -1:
+                    dist[ni] = d + 1
+                    q.append((nr, nc))
+    return dist.astype(np_cpu.float64)
+
+
+def assert_bfs_matches_pykokkos(N: int, M: int, mat, val, max_arr) -> None:
+    ref = reference_grid_bfs_distances(N, M, mat)
+    val_h = _view_to_numpy_host(val)
+    max_h = float(_view_to_numpy_host(max_arr)[0])
+    if not np_cpu.allclose(val_h, ref, rtol=0.0, atol=1e-9):
+        bad = np_cpu.where(np_cpu.abs(val_h - ref) > 1e-9)[0][:16]
+        raise AssertionError(
+            f"distance mismatch at linear indices (first few): {bad.tolist()}"
+        )
+    ref_max = float(np_cpu.max(ref))
+    if not np_cpu.isclose(max_h, ref_max, rtol=0.0, atol=1e-9):
+        raise AssertionError(f"max distance mismatch: got {max_h}, expected {ref_max}")
+    print("BFS correctness check: OK (matches queue-based NumPy reference)")
 
 
 def main(N: int, M: int):
@@ -29,6 +83,8 @@ def main(N: int, M: int):
     pk.parallel_for(N, extend2D, N=N, max_arr=max_arr, max_arr2D=max_arr2D)
     pk.parallel_for(N, reduce1D, N=N, max_arr=max_arr, max_arr2D=max_arr2D)
 
+    assert_bfs_matches_pykokkos(N, M, mat, val, max_arr)
+
     print(f"\ndistance of every cell:\n")
     for i in range(element):
         print(f"val ({val[i]})  ", end="")
@@ -73,21 +129,21 @@ def check_vis(
                 if min_val > val[i - M]:
                     min_val = val[i - M]
 
-            # check the neighbor on the next row
+        # check the neighbor on the next row
         if i // M < (N - 1):
             if visited[i + M] == 1:
                 flag = 1
                 if min_val > val[i + M]:
                     min_val = val[i + M]
 
-            # check the neighbor on the left
+        # check the neighbor on the left
         if i % M > 0:
             if visited[i - 1] == 1:
                 flag = 1
                 if min_val > val[i - 1]:
                     min_val = val[i - 1]
 
-            # check the neighbor on the right
+        # check the neighbor on the right
         if i % M < (M - 1):
             if visited[i + 1] == 1:
                 flag = 1
@@ -96,10 +152,10 @@ def check_vis(
 
         # if there is at least one neighbor visited, the value of
         # the current index can be updated and should be marked as visited
-    if flag == 1:
-        if val[i] > min_val:
-            val[i] = min_val + 1
-        visited[i] = 1
+        if flag == 1:
+            if val[i] > min_val:
+                val[i] = min_val + 1
+            visited[i] = 1
 
     ################################
     # findmax will find the maximum value of cell in each row

examples/pykokkos/binsort.py

@@ -12,11 +12,14 @@ def main():
     view = np.zeros(total_threads, dtype=np.int32)
 
     pk.parallel_for(total_threads, work, total_threads=total_threads, view=view)
+    max_bins = total_threads // 2
+    min_key = total_threads
+    max_key = total_threads * 2 - 1
     bin_op = pk.BinOp1D(
         view,
-        (total_threads // 2),
-        total_threads,
-        total_threads * 2 - 1,
+        max_bins,
+        min_key,
+        max_key,
     )
     bin_sort = pk.BinSort(view, bin_op)
     bin_sort.create_permute_vector()
@@ -25,10 +28,31 @@ def main():
     bin_count = bin_sort.get_bin_count()
     bin_sort.sort(view)
 
-    print(view)
-    print(permute_vector)
-    print(bin_offsets)
-    print(bin_count)
+    print(
+        "PyKokkos BinSort demo: fill a 1D key view on the device, bucket keys into "
+        f"{max_bins} bins over [{min_key}, {max_key}], then sort.\n"
+        f"  Initial keys: view[i] = i + {total_threads} (see work unit).\n"
+    )
+    print(
+        "Sorted keys (same 1D view, after bin_sort.sort(view) — in-place reorder):\n",
+        view,
+        "\n",
+    )
+    print(
+        "Permute vector from create_permute_vector / get_permute_vector — "
+        "indices describing how elements were reordered:\n",
+        permute_vector,
+        "\n",
+    )
+    print(
+        "Bin offsets (get_bin_offsets) — start index of each bin in the sorted layout:\n",
+        bin_offsets,
+        "\n",
+    )
+    print(
+        "Bin counts (get_bin_count) — number of keys in each bin:\n",
+        bin_count,
+    )
 
 
 @pk.workunit

examples/pykokkos/classtypes.py

@@ -11,13 +11,16 @@ def test(self) -> float:
 
 
 @pk.workunit
-def work(tid: int):
-    pk.printf("%d\n", tid)
+def work(tid: int, acc: pk.Acc[pk.double]) -> None:
+    tc: TestClass = TestClass(float(tid))
+    acc += tc.test()
 
 
 def main():
     total_threads: int = 10
-    pk.parallel_for(total_threads, work)
+    result: float = pk.parallel_reduce(total_threads, work)
+    expected: float = sum(2.0 * float(i) for i in range(total_threads))
+    print(f"parallel_reduce: {result} (expected {expected})")
 
 
 if __name__ == "__main__":

examples/pykokkos/subviews.py

@@ -1,4 +1,3 @@
-import numpy as np
 import pykokkos as pk
 
 if pk.get_default_space() in pk.DeviceExecutionSpace:
@@ -8,15 +7,24 @@
 
 
 def main():
-    view = np.zeros((10, 10), dtype=np.int32)
-    subview = view[3, 2:5]
-    pk.parallel_for(10, work, view=view)
+    n: int = 10
+    # 2D view; each thread takes a row subview and a column-range subview inside the workunit.
+    view = np.zeros((n, n), dtype=np.int32)
+    pk.parallel_for(n, work, view=view)
+    print(
+        "PyKokkos subviews: each iteration builds Kokkos::subview row = view[i, :] and "
+        "band = view[i, 2:5], then writes through those 1D subviews (not view[i][j] alone).\n"
+    )
     print(view)
 
 
 @pk.workunit
 def work(i: int, view: pk.View2D[pk.int32]):
-    view[i][i] = 1
+    # Plain `name = view[...]` becomes Kokkos::subview (annotated assign does not).
+    row = view[i, :]
+    band = view[i, 2:5]
+    row[i] = 1
+    band[i % 3] = i + 1
 
 
 if __name__ == "__main__":