remove inclusive_scan from other branch

Author: IvanGrigorik

examples/kokkos/inclusive_scan_team.py

@@ -0,0 +1,75 @@
+import pykokkos as pk
+
+
+@pk.workunit
+def init_data(i: int, view: pk.View1D[int]):
+    view[i] = i + 1
+
+
+# Test lower_bound with scratch memory
+@pk.workunit
+def team_lower_bound(team_member: pk.TeamMember, view: pk.View1D[int], result_view: pk.View1D[int]):
+    team_size: int = team_member.team_size()
+    offset: int = team_member.league_rank() * team_size
+    localIdx: int = team_member.team_rank()
+    globalIdx: int = offset + localIdx
+    team_rank: int = team_member.team_rank()
+
+    # Allocate scratch memory for sorted data
+    scratch: pk.ScratchView1D[int] = pk.ScratchView1D(
+        team_member.team_scratch(0), team_size
+    )
+
+    # Copy data to scratch and make it sorted within the team
+    scratch[team_rank] = view[globalIdx]
+    team_member.team_barrier()
+
+    # Now use lower_bound to find position in scratch
+    # For example, find lower bound for the value at team_rank position
+    search_value: int = team_rank * 2  # Search for a value
+    
+    # Find lower bound in scratch memory
+    bound_idx: int = pk.lower_bound(scratch, team_size, search_value)
+    
+    # Store result
+    result_view[globalIdx] = bound_idx
+
+
+# Test lower_bound with regular view
+@pk.workunit
+def test_lower_bound_view(i: int, view: pk.View1D[int], result_view: pk.View1D[int]):
+    # Find lower bound for value i in the first 10 elements
+    search_value: int = i
+    bound_idx: int = pk.lower_bound(view, 10, search_value)
+    result_view[i] = bound_idx
+
+
+def main():
+    N = 64
+    team_size = 32
+    num_teams = (N + team_size - 1) // team_size
+
+    view: pk.View1D[int] = pk.View([N], int)
+    result_view: pk.View1D[int] = pk.View([N], int)
+    
+    p_init = pk.RangePolicy(pk.ExecutionSpace.OpenMP, 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}")
+    print(f"Initial view: {view}")
+
+    # Test with TeamPolicy (scratch memory)
+    team_policy = pk.TeamPolicy(pk.ExecutionSpace.OpenMP, num_teams, team_size)
+    
+    print("\nRunning lower_bound with scratch memory...")
+    pk.parallel_for(team_policy, team_lower_bound, view=view, result_view=result_view)
+    print(f"Result (scratch lower_bound): {result_view}")
+
+    # Test with RangePolicy (regular view)
+    print("\nRunning lower_bound with regular view...")
+    pk.parallel_for(p_init, test_lower_bound_view, view=view, result_view=result_view)
+    print(f"Result (view lower_bound): {result_view}")
+
+
+if __name__ == "__main__":
+    main()

examples/kokkos/lower_bound_test.py

@@ -0,0 +1,75 @@
+import pykokkos as pk
+
+
+@pk.workunit
+def init_data(i: int, view: pk.View1D[int]):
+    view[i] = i + 1
+
+
+# Test upper_bound with scratch memory
+@pk.workunit
+def team_upper_bound(team_member: pk.TeamMember, view: pk.View1D[int], result_view: pk.View1D[int]):
+    team_size: int = team_member.team_size()
+    offset: int = team_member.league_rank() * team_size
+    localIdx: int = team_member.team_rank()
+    globalIdx: int = offset + localIdx
+    team_rank: int = team_member.team_rank()
+
+    # Allocate scratch memory for sorted data
+    scratch: pk.ScratchView1D[int] = pk.ScratchView1D(
+        team_member.team_scratch(0), team_size
+    )
+
+    # Copy data to scratch and make it sorted within the team
+    scratch[team_rank] = view[globalIdx]
+    team_member.team_barrier()
+
+    # Now use upper_bound to find position in scratch
+    # For example, find upper bound for the value at team_rank position
+    search_value: int = team_rank * 2  # Search for a value
+    
+    # Find upper bound in scratch memory
+    bound_idx: int = pk.upper_bound(scratch, team_size, search_value)
+    
+    # Store result
+    result_view[globalIdx] = bound_idx
+
+
+# Test upper_bound with regular view
+@pk.workunit
+def test_upper_bound_view(i: int, view: pk.View1D[int], result_view: pk.View1D[int]):
+    # Find upper bound for value i in the first 10 elements
+    search_value: int = i
+    bound_idx: int = pk.upper_bound(view, 10, search_value)
+    result_view[i] = bound_idx
+
+
+def main():
+    N = 64
+    team_size = 32
+    num_teams = (N + team_size - 1) // team_size
+
+    view: pk.View1D[int] = pk.View([N], int)
+    result_view: pk.View1D[int] = pk.View([N], int)
+    
+    p_init = pk.RangePolicy(pk.ExecutionSpace.Cuda, 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}")
+    print(f"Initial view: {view}")
+
+    # Test with TeamPolicy (scratch memory)
+    team_policy = pk.TeamPolicy(pk.ExecutionSpace.Cuda, num_teams, team_size)
+    
+    print("\nRunning upper_bound with scratch memory...")
+    pk.parallel_for(team_policy, team_upper_bound, view=view, result_view=result_view)
+    print(f"Result (scratch upper_bound): {result_view}")
+
+    # Test with RangePolicy (regular view)
+    print("\nRunning upper_bound with regular view...")
+    pk.parallel_for(p_init, test_upper_bound_view, view=view, result_view=result_view)
+    print(f"Result (view upper_bound): {result_view}")
+
+
+if __name__ == "__main__":
+    main()

examples/kokkos/upper_bound_test.py

@@ -54,7 +54,7 @@ def __init__(self, members: PyKokkosMembers, pk_import: str, path: str):
         self.global_symbols.update(math_functions)
         self.global_symbols.update(allowed_types)
         self.global_symbols.update(view_dtypes)
-        self.global_symbols.update(["self", "range", "math", "List", "abs", "inclusive_scan", "upper_bound"])
+        self.global_symbols.update(["self", "range", "math", "List", "abs", "upper_bound", "lower_bound"])
         self.global_symbols.add(pk_import)
 
         self.global_symbols.update([field.declname for field in members.fields])

pykokkos/core/translators/symbols_pass.py

@@ -318,51 +318,6 @@ def visit_Call(self, node: ast.Call) -> cppast.CallExpr:
 
             return real_number_call
 
-        # Kokkos `inclusive_scan`
-        if name == "inclusive_scan":
-            # Check if it's called via pk.inclusive_scan
-            is_pk_call = (
-                isinstance(node.func, ast.Attribute)
-                and isinstance(node.func.value, ast.Name)
-                and node.func.value.id == self.pk_import
-            )
-
-            if not is_pk_call:
-                return super().visit_Call(node)
-
-            # Expected signature: pk.inclusive_scan(team_member, view, length)
-            # or: pk.inclusive_scan(team_member, view) - uses view size
-            if len(args) < 2 or len(args) > 3:
-                self.error(
-                    node,
-                    "pk.inclusive_scan() takes 2 or 3 arguments: team_member, view, [length]",
-                )
-
-            team_member_expr = args[0]
-            view_expr = args[1]
-
-            # Create Kokkos::Experimental::begin() and end() calls
-            begin_function = cppast.DeclRefExpr("Kokkos::Experimental::begin")
-            end_function = cppast.DeclRefExpr("Kokkos::Experimental::end")
-            view_begin = cppast.CallExpr(begin_function, [view_expr])
-
-            if len(args) == 3:
-                # Use provided length: begin + length
-                length_expr = args[2]
-                view_begin_for_end = cppast.CallExpr(begin_function, [view_expr])
-                view_end = cppast.BinaryOperator(
-                    view_begin_for_end, length_expr, cppast.BinaryOperatorKind.Add
-                )
-            else:
-                # Use end() when no length is provided
-                view_end = cppast.CallExpr(end_function, [view_expr])
-
-            # Create Kokkos::Experimental::inclusive_scan call
-            function = cppast.DeclRefExpr("Kokkos::Experimental::inclusive_scan")
-            scan_args = [team_member_expr, view_begin, view_end, view_begin]
-
-            return cppast.CallExpr(function, scan_args)
-
         # Custom `upper_bound` implementation using binary search
         if name == "upper_bound":
             # Check if it's called via pk.upper_bound

pykokkos/core/visitors/workunit_visitor.py

@@ -1,5 +1,5 @@
 from .accumulator import Acc
-from .algorithms import inclusive_scan, lower_bound, upper_bound
+from .algorithms import lower_bound, upper_bound
 from .atomic.atomic_fetch_op import (
     atomic_fetch_add, atomic_fetch_and, atomic_fetch_div,
     atomic_fetch_lshift, atomic_fetch_max, atomic_fetch_min,

pykokkos/interface/__init__.py

@@ -1,5 +1,4 @@
-from .inclusive_scan import inclusive_scan
 from .lower_bound import lower_bound
 from .upper_bound import upper_bound
 
-__all__ = ["inclusive_scan", "lower_bound", "upper_bound"]
+__all__ = ["lower_bound", "upper_bound"]