add lower and upper bound functions, using binary search

Author: IvanGrigorik

pykokkos/core/translators/symbols_pass.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"])
+        self.global_symbols.update(["self", "range", "math", "List", "abs", "inclusive_scan", "upper_bound"])
         self.global_symbols.add(pk_import)
 
         self.global_symbols.update([field.declname for field in members.fields])

pykokkos/core/visitors/workunit_visitor.py

@@ -363,6 +363,76 @@ def visit_Call(self, node: ast.Call) -> cppast.CallExpr:
 
             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
+            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.upper_bound(view, size, value)
+            if len(args) != 3:
+                self.error(
+                    node,
+                    "pk.upper_bound() takes 3 arguments: view, size, value",
+                )
+
+            view_expr = args[0]
+            size_expr = args[1]
+            value_expr = args[2]
+
+            # Generate binary search lambda inline
+            from pykokkos.interface.algorithms.upper_bound import generate_upper_bound_binary_search
+            
+            # Create lambda body with binary search
+            lambda_body = generate_upper_bound_binary_search(view_expr, size_expr, value_expr)
+            
+            # Create and invoke lambda
+            lambda_expr = cppast.LambdaExpr("[&]", [], lambda_body)
+            lambda_call = cppast.CallExpr(lambda_expr, [])
+
+            return lambda_call
+
+        # Custom `lower_bound` implementation using binary search
+        if name == "lower_bound":
+            # Check if it's called via pk.lower_bound
+            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.lower_bound(view, size, value)
+            if len(args) != 3:
+                self.error(
+                    node,
+                    "pk.lower_bound() takes 3 arguments: view, size, value",
+                )
+
+            view_expr = args[0]
+            size_expr = args[1]
+            value_expr = args[2]
+
+            # Generate binary search lambda inline
+            from pykokkos.interface.algorithms.lower_bound import generate_lower_bound_binary_search
+            
+            # Create lambda body with binary search
+            lambda_body = generate_lower_bound_binary_search(view_expr, size_expr, value_expr)
+            
+            # Create and invoke lambda
+            lambda_expr = cppast.LambdaExpr("[&]", [], lambda_body)
+            lambda_call = cppast.CallExpr(lambda_expr, [])
+
+            return lambda_call
+
         return super().visit_Call(node)
 
     def is_nested_call(self, node: ast.FunctionDef) -> bool:

pykokkos/interface/__init__.py

@@ -1,5 +1,5 @@
 from .accumulator import Acc
-from .algorithms.inclusive_scan import inclusive_scan
+from .algorithms import inclusive_scan, 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/algorithms/__init__.py

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

pykokkos/interface/algorithms/lower_bound.py

@@ -0,0 +1,90 @@
+from pykokkos.interface.views import ViewType
+from pykokkos.core import cppast
+
+def lower_bound(view: ViewType, size: int, value) -> int:
+    """
+    Perform a lower bound search on a view
+    
+    Returns the index of the first element not less than (i.e. greater or equal to) value,
+    similar to std::lower_bound or thrust::lower_bound.
+    
+    :param view: the view to search (must be sorted)
+    :param size: the number of elements to search
+    :param value: the value to search for
+    :returns: the index of the first element >= value
+    """
+    pass
+
+
+def generate_lower_bound_binary_search(
+    view_expr: cppast.Expr, size_expr: cppast.Expr, value_expr: cppast.Expr
+) -> cppast.CompoundStmt:
+    """
+    Generate binary search implementation for lower_bound.
+    Returns a CompoundStmt that implements:
+    
+    int left = 0;
+    int right = size;
+    int mid;
+    while (left < right) {
+        mid = left + (right - left) / 2;
+        if (view[mid] < value) {
+            left = mid + 1;
+        } else {
+            right = mid;
+        }
+    }
+    return left;
+    """
+    
+    # Variable declarations
+    int_type = cppast.PrimitiveType("int32_t")
+    
+    # int left = 0;
+    left_var = cppast.DeclRefExpr("left")
+    left_init = cppast.IntegerLiteral(0)
+    left_decl = cppast.VarDecl(int_type, left_var, left_init)
+    left_stmt = cppast.DeclStmt(left_decl)
+    
+    # int right = size;
+    right_var = cppast.DeclRefExpr("right")
+    right_decl = cppast.VarDecl(int_type, right_var, size_expr)
+    right_stmt = cppast.DeclStmt(right_decl)
+    
+    # int mid;
+    mid_var = cppast.DeclRefExpr("mid")
+    mid_decl = cppast.VarDecl(int_type, mid_var, None)
+    mid_stmt = cppast.DeclStmt(mid_decl)
+    
+    # while (left < right)
+    while_cond = cppast.BinaryOperator(left_var, right_var, cppast.BinaryOperatorKind.LT)
+    
+    # mid = left + (right - left) / 2;
+    right_minus_left = cppast.BinaryOperator(right_var, left_var, cppast.BinaryOperatorKind.Sub)
+    div_expr = cppast.BinaryOperator(right_minus_left, cppast.IntegerLiteral(2), cppast.BinaryOperatorKind.Div)
+    mid_calc = cppast.BinaryOperator(left_var, div_expr, cppast.BinaryOperatorKind.Add)
+    mid_assign = cppast.AssignOperator([mid_var], mid_calc, cppast.BinaryOperatorKind.Assign)
+    
+    # if (view[mid] < value)
+    view_access = cppast.CallExpr(view_expr, [mid_var])
+    if_cond = cppast.BinaryOperator(view_access, value_expr, cppast.BinaryOperatorKind.LT)
+    
+    # left = mid + 1;
+    mid_plus_one = cppast.BinaryOperator(mid_var, cppast.IntegerLiteral(1), cppast.BinaryOperatorKind.Add)
+    left_assign = cppast.AssignOperator([left_var], mid_plus_one, cppast.BinaryOperatorKind.Assign)
+    
+    # right = mid;
+    right_assign = cppast.AssignOperator([right_var], mid_var, cppast.BinaryOperatorKind.Assign)
+    
+    # if-else statement
+    if_stmt = cppast.IfStmt(if_cond, left_assign, right_assign)
+    
+    # while body
+    while_body = cppast.CompoundStmt([mid_assign, if_stmt])
+    while_stmt = cppast.WhileStmt(while_cond, while_body)
+    
+    # return left;
+    return_stmt = cppast.ReturnStmt(left_var)
+    
+    # Complete function body
+    return cppast.CompoundStmt([left_stmt, right_stmt, mid_stmt, while_stmt, return_stmt])

pykokkos/interface/algorithms/upper_bound.py

@@ -0,0 +1,90 @@
+from pykokkos.interface.views import ViewType
+from pykokkos.core import cppast
+
+def upper_bound(view: ViewType, size: int, value) -> int:
+    """
+    Perform an upper bound search on a view
+    
+    Returns the index of the first element greater than value,
+    similar to std::upper_bound or thrust::upper_bound.
+    
+    :param view: the view to search (must be sorted)
+    :param size: the number of elements to search
+    :param value: the value to search for
+    :returns: the index of the first element greater than value
+    """
+    pass
+
+
+def generate_upper_bound_binary_search(
+    view_expr: cppast.Expr, size_expr: cppast.Expr, value_expr: cppast.Expr
+) -> cppast.CompoundStmt:
+    """
+    Generate binary search implementation for upper_bound.
+    Returns a CompoundStmt that implements:
+    
+    int left = -1;
+    int right = size;
+    int mid;
+    while (left + 1 < right) {
+        mid = left + ((right - left) >> 1);
+        if (view[mid] > value) {
+            right = mid;
+        } else {
+            left = mid;
+        }
+    }
+    return right;
+    """
+    
+    # Variable declarations
+    int_type = cppast.PrimitiveType("int32_t")
+    
+    # int left = 0;
+    left_var = cppast.DeclRefExpr("left")
+    left_init = cppast.IntegerLiteral(0)
+    left_decl = cppast.VarDecl(int_type, left_var, left_init)
+    left_stmt = cppast.DeclStmt(left_decl)
+    
+    # int right = size;
+    right_var = cppast.DeclRefExpr("right")
+    right_decl = cppast.VarDecl(int_type, right_var, size_expr)
+    right_stmt = cppast.DeclStmt(right_decl)
+    
+    # int mid;
+    mid_var = cppast.DeclRefExpr("mid")
+    mid_decl = cppast.VarDecl(int_type, mid_var, None)
+    mid_stmt = cppast.DeclStmt(mid_decl)
+    
+    # while (left < right)
+    while_cond = cppast.BinaryOperator(left_var, right_var, cppast.BinaryOperatorKind.LT)
+    
+    # mid = left + (right - left) / 2;
+    right_minus_left = cppast.BinaryOperator(right_var, left_var, cppast.BinaryOperatorKind.Sub)
+    div_expr = cppast.BinaryOperator(right_minus_left, cppast.IntegerLiteral(2), cppast.BinaryOperatorKind.Div)
+    mid_calc = cppast.BinaryOperator(left_var, div_expr, cppast.BinaryOperatorKind.Add)
+    mid_assign = cppast.AssignOperator([mid_var], mid_calc, cppast.BinaryOperatorKind.Assign)
+    
+    # if (view[mid] > value)
+    view_access = cppast.CallExpr(view_expr, [mid_var])
+    if_cond = cppast.BinaryOperator(view_access, value_expr, cppast.BinaryOperatorKind.GT)
+    
+    # right = mid;
+    right_assign = cppast.AssignOperator([right_var], mid_var, cppast.BinaryOperatorKind.Assign)
+    
+    # left = mid + 1;
+    mid_plus_one = cppast.BinaryOperator(mid_var, cppast.IntegerLiteral(1), cppast.BinaryOperatorKind.Add)
+    left_assign = cppast.AssignOperator([left_var], mid_plus_one, cppast.BinaryOperatorKind.Assign)
+    
+    # if-else statement
+    if_stmt = cppast.IfStmt(if_cond, right_assign, left_assign)
+    
+    # while body
+    while_body = cppast.CompoundStmt([mid_assign, if_stmt])
+    while_stmt = cppast.WhileStmt(while_cond, while_body)
+    
+    # return left;
+    return_stmt = cppast.ReturnStmt(left_var)
+    
+    # Complete function body
+    return cppast.CompoundStmt([left_stmt, right_stmt, mid_stmt, while_stmt, return_stmt])