fix pykokkos to be able to run ExaMiniMD under Debug mode in new repo

Author: IvanGrigorik

pykokkos/pykokkos/core/run_debug.py

@@ -122,7 +122,7 @@ def run_workunit_debug(
 
     elif isinstance(policy, TeamThreadRange) or isinstance(policy, ThreadVectorRange):
         for i in range(policy.count):
-            call_workunit(operation, workunit, TeamMember(i, 0), acc, **kwargs)
+            call_workunit(operation, workunit, i, acc, **kwargs)
 
     else:
         if isinstance(policy, MDRangePolicy):

pykokkos/pykokkos/core/runtime.py

@@ -418,6 +418,9 @@ def is_debug(self, space: ExecutionSpace) -> bool:
         :returns: True or False
         """
 
+        if km.get_default_space() is ExecutionSpace.Debug:
+            return True
+
         return space is ExecutionSpace.Debug or (
             space is ExecutionSpace.Default
             and km.get_default_space() is ExecutionSpace.Debug

pykokkos/pykokkos/interface/accumulator.py

@@ -5,56 +5,70 @@ class Acc(Generic[TypeVar("T")]):
     def __init__(self, val):
         self.val = val
 
+    # Non-augmented arithmetic (`n + 1`, `n * 2`, ...) must not mutate `n`,
+    # matching normal Python operator semantics.
     def __add__(self, other):
-        self.val = self.val + other
-        return self
+        return self.val + other
 
     def __radd__(self, other):
-        self.val = self.val + other
-        return self
+        return other + self.val
 
     def __sub__(self, other):
-        self.val = self.val - other
-        return self
+        return self.val - other
 
     def __rsub__(self, other):
-        self.val = other - self.val
-        return self
+        return other - self.val
 
     def __mul__(self, other):
-        self.val = self.val * other
-        return self
+        return self.val * other
 
     def __rmul__(self, other):
-        self.val = self.val * other
-        return self
+        return other * self.val
 
     def __truediv__(self, other):
-        self.val = self.val / other
-        return self
+        return self.val / other
 
     def __rtruediv__(self, other):
-        self.val = other / self.val
-        return self
+        return other / self.val
 
     def __floordiv__(self, other):
-        self.val = self.val // other
-        return self
+        return self.val // other
 
     def __rfloordiv__(self, other):
-        self.val = other // self.val
-        return self
+        return other // self.val
 
     def __mod__(self, other):
-        self.val = self.val % other
+        return self.val % other
+
+    def __rmod__(self, other):
+        return other % self.val
+
+    def __neg__(self):
+        return -self.val
+
+    # Augmented assignment (`n += 1`, ...) is where mutation belongs.
+    def __iadd__(self, other):
+        self.val = self.val + other
         return self
 
-    def __rmov__(self, other):
-        self.val = other % self.val
+    def __isub__(self, other):
+        self.val = self.val - other
         return self
 
-    def __neg__(self):
-        self.val = -self.val
+    def __imul__(self, other):
+        self.val = self.val * other
+        return self
+
+    def __itruediv__(self, other):
+        self.val = self.val / other
+        return self
+
+    def __ifloordiv__(self, other):
+        self.val = self.val // other
+        return self
+
+    def __imod__(self, other):
+        self.val = self.val % other
         return self
 
     def __index__(self):

pykokkos/pykokkos/interface/atomic/atomic_fetch_op.py

@@ -11,28 +11,32 @@ def atomic_fetch_add(
     view: View, indices: List[int], value: Union[int, float]
 ) -> Union[int, float]:
     inner = reduce(operator.getitem, indices[:-1], view)
+    old_result = inner[indices[-1]]
     inner[indices[-1]] += value
-    return inner[indices[-1]]
+    return old_result
 
 
 def atomic_fetch_and(view: View, indices: List[int], value: int) -> int:
     inner = reduce(operator.getitem, indices[:-1], view)
+    old_result = inner[indices[-1]]
     inner[indices[-1]] &= value
-    return inner[indices[-1]]
+    return old_result
 
 
 def atomic_fetch_div(
     view: View, indices: List[int], value: Union[int, float]
 ) -> Union[int, float]:
     inner = reduce(operator.getitem, indices[:-1], view)
+    old_result = inner[indices[-1]]
     inner[indices[-1]] /= value
-    return inner[indices[-1]]
+    return old_result
 
 
 def atomic_fetch_lshift(view: View, indices: List[int], value: int) -> int:
     inner = reduce(operator.getitem, indices[:-1], view)
+    old_result = inner[indices[-1]]
     inner[indices[-1]] <<= value
-    return inner[indices[-1]]
+    return old_result
 
 
 def atomic_fetch_max(
@@ -55,42 +59,48 @@ def atomic_fetch_min(
 
 def atomic_fetch_mod(view: View, indices: List[int], value: int) -> int:
     inner = reduce(operator.getitem, indices[:-1], view)
+    old_result = inner[indices[-1]]
     inner[indices[-1]] %= value
-    return inner[indices[-1]]
+    return old_result
 
 
 def atomic_fetch_mul(
     view: View, indices: List[int], value: Union[int, float]
 ) -> Union[int, float]:
     inner = reduce(operator.getitem, indices[:-1], view)
+    old_result = inner[indices[-1]]
     inner[indices[-1]] *= value
-    return inner[indices[-1]]
+    return old_result
 
 
 def atomic_fetch_or(view: View, indices: List[int], value: int) -> int:
     inner = reduce(operator.getitem, indices[:-1], view)
+    old_result = inner[indices[-1]]
     inner[indices[-1]] |= value
-    return inner[indices[-1]]
+    return old_result
 
 
 def atomic_fetch_rshift(view: View, indices: List[int], value: int) -> int:
     inner = reduce(operator.getitem, indices[:-1], view)
+    old_result = inner[indices[-1]]
     inner[indices[-1]] >>= value
-    return inner[indices[-1]]
+    return old_result
 
 
 def atomic_fetch_sub(
     view: View, indices: List[int], value: Union[int, float]
 ) -> Union[int, float]:
     inner = reduce(operator.getitem, indices[:-1], view)
+    old_result = inner[indices[-1]]
     inner[indices[-1]] -= value
-    return inner[indices[-1]]
+    return old_result
 
 
 def atomic_fetch_xor(view: View, indices: List[int], value: int) -> int:
     inner = reduce(operator.getitem, indices[:-1], view)
+    old_result = inner[indices[-1]]
     inner[indices[-1]] ^= value
-    return inner[indices[-1]]
+    return old_result
 
 
 def atomic_compare_exchange(

pykokkos/pykokkos/interface/bin_sort.py

@@ -1,6 +1,8 @@
-from typing import List, Union
+from typing import List, Optional, Union
 
-from .views import View
+import numpy as np
+
+from .views import View, ViewType, array
 
 
 class BinOp:
@@ -20,11 +22,31 @@ def __init__(
     def get_type(dim: int, key_view_type: str) -> str:
         return f"Kokkos::BinOp{dim}D<{key_view_type}>"
 
+    def _keys_array(self) -> np.ndarray:
+        return self.keys.data if isinstance(self.keys, ViewType) else np.asarray(self.keys)
+
+    def num_bins(self) -> int:
+        raise NotImplementedError
+
+    def bin_indices(self) -> np.ndarray:
+        """Flat bin index per key, used by the interpreted (Debug) execution path."""
+        raise NotImplementedError
+
 
 class BinOp1D(BinOp):
     def __init__(self, keys: View, max_bins: int, min_value: float, max_value: float):
         super().__init__(keys, max_bins, min_value, max_value)
 
+    def num_bins(self) -> int:
+        return self.max_bins
+
+    def bin_indices(self) -> np.ndarray:
+        keys_arr = self._keys_array()
+        span = self.max_value - self.min_value
+        scale = self.max_bins / span if span > 0 else 0.0
+        idx = np.floor((keys_arr - self.min_value) * scale).astype(np.int64)
+        return np.clip(idx, 0, self.max_bins - 1)
+
 
 class BinOp3D(BinOp):
     def __init__(
@@ -36,28 +58,66 @@ def __init__(
     ):
         super().__init__(keys, max_bins, min_value, max_value)
 
+    def num_bins(self) -> int:
+        nbx, nby, nbz = self.max_bins
+        return nbx * nby * nbz
+
+    def bin_indices(self) -> np.ndarray:
+        """Row-major flat index: ix * nby * nbz + iy * nbz + iz."""
+        keys_arr = self._keys_array()
+        nbx, nby, nbz = self.max_bins
+        idx3 = np.empty((keys_arr.shape[0], 3), dtype=np.int64)
+        for d, n in enumerate((nbx, nby, nbz)):
+            span = self.max_value[d] - self.min_value[d]
+            scale = n / span if span > 0 else 0.0
+            col = np.floor((keys_arr[:, d] - self.min_value[d]) * scale).astype(np.int64)
+            idx3[:, d] = np.clip(col, 0, n - 1)
+        return idx3[:, 0] * (nby * nbz) + idx3[:, 1] * nbz + idx3[:, 2]
+
 
 class BinSort:
     def __init__(self, keys: View, bin_op: BinOp, sort_within_bins: bool = False):
         self.keys = keys
         self.bin_op = bin_op
         self.sort_within_bins = sort_within_bins
+        self._permute_vector: Optional[np.ndarray] = None
+        self._bin_count: Optional[np.ndarray] = None
+        self._bin_offsets: Optional[np.ndarray] = None
 
     @staticmethod
     def get_type(key_view_type: str, bin_op_type: str, space: str) -> str:
         return f"Kokkos::BinSort<{key_view_type},{bin_op_type},{space},int>"
 
     def sort(self, values: View) -> None:
-        pass
+        if self._permute_vector is None:
+            raise RuntimeError("create_permute_vector() must be called before sort()")
+        data = values.data if isinstance(values, ViewType) else values
+        n = len(self._permute_vector)
+        data[:n] = data[:n][self._permute_vector]
 
     def get_bin_count(self) -> View:
-        pass
+        return array(self._bin_count)
 
     def get_bin_offsets(self) -> View:
-        pass
+        return array(self._bin_offsets)
 
     def get_permute_vector(self) -> View:
-        pass
+        return array(self._permute_vector)
 
     def create_permute_vector(self) -> None:
-        pass
+        """Counting sort by bin index (interpreted-execution fallback for Kokkos::BinSort)."""
+        bin_ids = self.bin_op.bin_indices()
+        n_bins = self.bin_op.num_bins()
+
+        counts = np.bincount(bin_ids, minlength=n_bins).astype(np.int32)
+        offsets = np.zeros(n_bins, dtype=np.int32)
+        if n_bins > 1:
+            offsets[1:] = np.cumsum(counts)[:-1]
+
+        # Stable sort so ties (same bin) keep their original relative order,
+        # matching sort_within_bins=False semantics.
+        order = np.argsort(bin_ids, kind="stable").astype(np.int32)
+
+        self._bin_count = counts
+        self._bin_offsets = offsets
+        self._permute_vector = order

pykokkos/pykokkos/interface/execution_policy.py

@@ -145,7 +145,12 @@ def __init__(self, *args):
         unpacked: Tuple = tuple(args)
 
         if len(unpacked) == 2:
-            space = km.get_execution_space_instance(km.get_default_space())
+            default_space = km.get_default_space()
+            space = (
+                default_space
+                if default_space is ExecutionSpace.Debug
+                else km.get_execution_space_instance(default_space)
+            )
             league_size = unpacked[0]
             team_size = unpacked[1]
             vector_length = -1
@@ -163,7 +168,12 @@ def __init__(self, *args):
                 team_size = third
                 vector_length = -1
             else:
-                space = km.get_execution_space_instance(km.get_default_space())
+                default_space = km.get_default_space()
+                space = (
+                    default_space
+                    if default_space is ExecutionSpace.Debug
+                    else km.get_execution_space_instance(default_space)
+                )
                 league_size = first
                 team_size = second
                 vector_length = third
@@ -189,7 +199,8 @@ def __init__(self, *args):
         if isinstance(space, ExecutionSpace):
             if space is ExecutionSpace.Default:
                 space = km.get_default_space()
-            space = ExecutionSpaceInstance(space)
+            if space is not ExecutionSpace.Debug:
+                space = ExecutionSpaceInstance(space)
 
         elif not isinstance(space, ExecutionSpaceInstance):
             raise TypeError(f"Invalid space argument {space}")

pykokkos/pykokkos/interface/parallel_dispatch.py

@@ -8,7 +8,7 @@
 from pykokkos.core.cppast import BuiltinType
 
 from .execution_policy import ExecutionPolicy, RangePolicy
-from .execution_space import ExecutionSpace, DeviceExecutionSpace
+from .execution_space import ExecutionSpace, ExecutionSpaceInstance, DeviceExecutionSpace
 from .views import ViewType, array
 
 from .interface_util import generic_error, get_filename, get_lineno
@@ -195,6 +195,11 @@ def check_workunit(workunit: Any) -> None:
         raise TypeError(f"ERROR: {workunit} is not a valid workunit")
 
 
+def policy_execution_space(policy: ExecutionPolicy) -> ExecutionSpace:
+    space = policy.space
+    return space.space if isinstance(space, ExecutionSpaceInstance) else space
+
+
 def convert_arrays(kwargs: Dict[str, Any], workunit: Callable, execution_space) -> None:
     """
     Convert all numpy, cupy and pytorch ndarray objects into pk Views
@@ -303,7 +308,7 @@ def parallel_for(*args, **kwargs) -> None:
 
     kwargs = dict(kwargs)
     handled_args: HandledArgs = handle_args(True, args)
-    convert_arrays(kwargs, handled_args.workunit, handled_args.policy.space.space)
+    convert_arrays(kwargs, handled_args.workunit, policy_execution_space(handled_args.policy))
 
     runtime_singleton.runtime.run_workunit(
         handled_args.name, handled_args.policy, handled_args.workunit, "for", **kwargs
@@ -320,7 +325,7 @@ def reduce_body(operation: str, *args, **kwargs) -> Union[float, int]:
 
     kwargs = dict(kwargs)
     handled_args: HandledArgs = handle_args(True, args)
-    convert_arrays(kwargs, handled_args.workunit, handled_args.policy.space.space)
+    convert_arrays(kwargs, handled_args.workunit, policy_execution_space(handled_args.policy))
 
     args_to_hash: List = []
     args_not_to_hash: Dict = {}