Update co oc kernel (#347)

Author: Intron7

src/rapids_singlecell/squidpy_gpu/_co_oc.py

@@ -1,6 +1,5 @@
 from __future__ import annotations
 
-import math
 from typing import TYPE_CHECKING
 
 import cupy as cp
@@ -9,9 +8,10 @@
 
 from ._utils import _assert_categorical_obs, _assert_spatial_basis
 from .kernels._co_oc import (
-    occur_count_kernel,
-    occur_count_kernel2,
-    occur_count_kernel_fast,
+    occur_count_kernel_pairwise,
+    occur_count_kernel_pairwise_fast,
+    occur_reduction_kernel_global,
+    occur_reduction_kernel_shared,
 )
 
 if TYPE_CHECKING:
@@ -99,6 +99,26 @@ def _find_min_max(spatial: NDArrayC) -> tuple[float, float]:
     return thres_min.astype(np.float32), thres_max.astype(np.float32)
 
 
+def calculate_optimal_k(target_occupancy=0.4):
+    props = cp.cuda.runtime.getDeviceProperties(0)
+
+    # Get key SM properties
+    shared_mem_per_sm = props["sharedMemPerMultiprocessor"]  # bytes
+    max_warps_per_sm = props["maxThreadsPerMultiProcessor"] // 32
+    block_size = 128  # Your current block size
+    warps_per_block = block_size // 32  # 4 warps per block
+
+    # Target blocks per SM based on desired occupancy
+    target_blocks = int(max_warps_per_sm * target_occupancy) // warps_per_block
+
+    # Maximum shared memory per block to achieve target occupancy
+    max_shared_per_block = shared_mem_per_sm // target_blocks
+
+    # Calculate max k value
+    max_k = max_shared_per_block // (block_size * cp.dtype("float32").itemsize)
+    return max_k
+
+
 def _co_occurrence_helper(
     spatial: NDArrayC, v_radium: NDArrayC, labs: NDArrayC, fast: bool = True
 ) -> NDArrayC:
@@ -125,35 +145,57 @@ def _co_occurrence_helper(
     k = len(labs_unique)
     l_val = len(v_radium) - 1
     thresholds = (v_radium[1:]) ** 2
-    shared_mem_size_fast = (k * k * 32) * cp.dtype("float32").itemsize
-    props = cp.cuda.runtime.getDeviceProperties(0)
-    # if the shared memory is enough, use the fast kernel
-    if props["sharedMemPerBlock"] > shared_mem_size_fast and fast:
-        counts = cp.zeros((l_val, k, k), dtype=cp.int32)
-        grid = (int(math.ceil(n / 32)), l_val)
-        block = (32, 1)
-        occur_count_kernel_fast(
-            grid,
-            block,
-            (spatial, thresholds, labs, counts, n, k, l_val),
-            shared_mem=shared_mem_size_fast,
-        )
-        reader = 1
-    # if the shared memory is not enough, use the slow kernel
-    else:
+    use_fast_kernel = False  # Flag to track which kernel path was taken
+    if fast:
+        # Optimize occupancy vs speed
+        can_use_fast_kernel = calculate_optimal_k(0.4) > k
+        # If shared memory is sufficient, use the fast kernel
+        if can_use_fast_kernel:
+            shared_mem_size_fast = (k * 128) * cp.dtype("float32").itemsize
+            counts = cp.zeros((l_val, k, k), dtype=cp.int32)
+            grid = (n, l_val)
+            block = (128, 1)
+            occur_count_kernel_pairwise_fast(
+                grid,
+                block,
+                (spatial, thresholds, labs, counts, n, k, l_val),
+                shared_mem=shared_mem_size_fast,
+            )
+            reader = 1
+            use_fast_kernel = True
+
+    # Fallback to the standard kernel if fast=False or shared memory was insufficient
+    if not use_fast_kernel:
         counts = cp.zeros((k, k, l_val * 2), dtype=cp.int32)
         grid = (n,)
         block = (32,)
-        occur_count_kernel(
+        occur_count_kernel_pairwise(
             grid, block, (spatial, thresholds, labs, counts, n, k, l_val)
         )
         reader = 0
 
     occ_prob = cp.empty((k, k, l_val), dtype=np.float32)
     shared_mem_size = (k * k + k) * cp.dtype("float32").itemsize
-    grid = (l_val,)
-    block = (32,)
-    occur_count_kernel2(
-        grid, block, (counts, occ_prob, k, l_val, reader), shared_mem=shared_mem_size
-    )
+    props = cp.cuda.runtime.getDeviceProperties(0)
+    if fast and shared_mem_size < props["sharedMemPerBlock"]:
+        grid2 = (l_val,)
+        block2 = (32,)
+        occur_reduction_kernel_shared(
+            grid2,
+            block2,
+            (counts, occ_prob, k, l_val, reader),
+            shared_mem=shared_mem_size,
+        )
+    else:
+        shared_mem_size = (k) * cp.dtype("float32").itemsize
+        grid2 = (l_val,)
+        block2 = (32,)
+        inter_out = cp.zeros((l_val, k, k), dtype=np.float32)
+        occur_reduction_kernel_global(
+            grid2,
+            block2,
+            (counts, inter_out, occ_prob, k, l_val, reader),
+            shared_mem=shared_mem_size,
+        )
+
     return occ_prob

src/rapids_singlecell/squidpy_gpu/kernels/_co_oc.py

@@ -2,9 +2,9 @@
 
 import cupy as cp
 
-kernel_code = r"""
+kernel_code_pairwise = r"""
 extern "C" __global__
-void occur_count_kernel(const float* __restrict__ spatial,
+void occur_count_kernel_pairwise(const float* __restrict__ spatial,
                         const float* __restrict__ thresholds,
                         const int* __restrict__ label_idx,
                         int* __restrict__ result,
@@ -54,11 +54,13 @@
     }
 }
 """
-occur_count_kernel = cp.RawKernel(kernel_code, "occur_count_kernel")
+occur_count_kernel_pairwise = cp.RawKernel(
+    kernel_code_pairwise, "occur_count_kernel_pairwise"
+)
 
-kernel_code_fast = r"""
+kernel_code_pairwise_fast = r"""
 extern "C" __global__
-void occur_count_kernel_fast(const float* __restrict__ spatial,
+void occur_count_kernel_pairwise_fast(const float* __restrict__ spatial,
                             const float* __restrict__ thresholds,
                             const int* __restrict__ label_idx,
                             int* __restrict__ result,
@@ -68,52 +70,52 @@
 {
     extern __shared__ float shared[];
     float* Y = shared;
-    int i = blockIdx.x * blockDim.x + threadIdx.x;
+    int i = blockIdx.x;
     int r = blockIdx.y;
-    for (int j = 0; j < k * k ; j++){
-        Y[k * k * threadIdx.x + j] = 0;
+    for (int j = threadIdx.x; j < k * blockDim.x ; j+= blockDim.x){
+        Y[j] = 0;
     }
     __syncthreads();
-    for(int j = i+1; j< n; j++){
-        float spx = spatial[i * 2];
-        float spy = spatial[i * 2 + 1];
+
+    float spx = spatial[i * 2];
+    int low = label_idx[i];
+    float spy = spatial[i * 2 + 1];
+    for(int j = i+1+ threadIdx.x; j< n; j+= blockDim.x){
         float dx = spx - spatial[j * 2];
         float dy = spy - spatial[j * 2 + 1];
         float dist_sq = dx * dx + dy * dy;
-        int low = label_idx[i];
         int high = label_idx[j];
-        if (high < low) {
-            int tmp = low;
-            low = high;
-            high = tmp;
-        }
-
         if (dist_sq <= thresholds[r]) {
-            int index = k * k * threadIdx.x + low * k + high;
+            int index = k * threadIdx.x + high;
             Y[index] += 1;
         }
     }
-
     __syncthreads();
-    for (int bin = threadIdx.x; bin < k * k; bin += blockDim.x) {
-        int blockSum = 0;
-        for (int t = 0; t < blockDim.x; t++) {
-            blockSum += Y[t * (k * k) + bin];
-        }
-        if (blockSum>0) atomicAdd(&result[r*(k*k)+bin], blockSum);
 
+    for (int j = threadIdx.x; j < k; j+= blockDim.x){
+        float sum = 0;
+        for (int t = 0; t < blockDim.x; t++){
+            int index = k * t + j;
+            sum += Y[index];
+        }
+        if (low < j){
+            if (sum>0) atomicAdd(&result[r*(k*k)+low*k+j], sum);
+        }
+        else{
+            if (sum>0) atomicAdd(&result[r*(k*k)+j*k+low], sum);
+        }
     }
     __syncthreads();
 }
 """
-
-# Compile the kernel.
-occur_count_kernel_fast = cp.RawKernel(kernel_code_fast, "occur_count_kernel_fast")
+occur_count_kernel_pairwise_fast = cp.RawKernel(
+    kernel_code_pairwise_fast, "occur_count_kernel_pairwise_fast"
+)
 
 
-kernel_code2 = r"""
+occur_reduction_kernel_code_shared = r"""
 extern "C" __global__
-void occur_reduction_kernel(const int* __restrict__ result,
+void occur_reduction_kernel_shared(const int* __restrict__ result,
                             float* __restrict__ out,
                             int k,
                             int l_val,
@@ -221,6 +223,122 @@
             out[i * (k * l_val) + j * l_val + r_th] = final_val;
         }
     }
+    __syncthreads();
+}
+"""
+occur_reduction_kernel_shared = cp.RawKernel(
+    occur_reduction_kernel_code_shared, "occur_reduction_kernel_shared"
+)
+
+occur_reduction_kernel_code_global = r"""
+extern "C" __global__
+void occur_reduction_kernel_global(const int* __restrict__ result,
+                            float* __restrict__ inter_out,
+                            float* __restrict__ out,
+                            int k,
+                            int l_val,
+                            int format)
+{
+    // Each block handles one threshold index.
+    int r_th = blockIdx.x;  // threshold index
+    if (r_th >= l_val)
+        return;
+    // Shared memory allocation
+    extern __shared__ float shared[];
+    float* Y = inter_out + r_th*k*k;
+    float* col_sum = shared;
+
+    int total_elements = k * k;
+
+    // --- Load counts for this threshold and convert to float---
+    if (format == 0){
+        for (int i = threadIdx.x; i < k; i += blockDim.x){
+            for (int j = 0; j<k;j++){
+                Y[i * k + j] += float(result[i * (k * l_val*2) + j * l_val*2 + r_th]);
+                Y[j * k + i] += float(result[i * (k * l_val*2) + j * l_val*2 + r_th]);
+                Y[i * k + j] += float(result[i * (k * l_val*2) + j * l_val*2 + r_th+l_val]);
+                Y[j * k + i] += float(result[i * (k * l_val*2) + j * l_val*2 + r_th+l_val]);
+            }
+        }
+    }
+    else{
+        for (int i = threadIdx.x; i < k; i += blockDim.x){
+            for (int j = 0; j<k;j++){
+                Y[i * k + j] += float(result[r_th * (k * k) + i * k + j]);
+                Y[j * k + i] += float(result[r_th * (k * k) + i * k + j]);
+            }
+        }
+    }
+    __syncthreads();
+
+    // Compute total sum of the counts
+    __shared__ float total;
+    float sum_val = 0.0f;
+    for (int idx = threadIdx.x; idx < total_elements; idx += blockDim.x) {
+        sum_val += Y[idx];
+    }
+    __syncthreads();
+    // Warp-level reduction
+    unsigned int mask = 0xFFFFFFFF;  // full warp mask
+    for (int offset = warpSize / 2; offset > 0; offset /= 2) {
+        sum_val += __shfl_down_sync(mask, sum_val, offset);
+    }
+    __syncthreads();
+    if (threadIdx.x == 0) {
+        total = sum_val;
+    }
+    __syncthreads();
+
+    // Normalize the matrix Y = Y / total (if total > 0)
+    if (total > 0.0f) {
+        for (int idx = threadIdx.x; idx < total_elements; idx += blockDim.x) {
+            Y[idx] = Y[idx] / total;
+        }
+    } else {
+        for (int i = threadIdx.x; i < k; i += blockDim.x) {
+            for (int j = 0; j < k; j++) {
+                out[i * (k * l_val) + j * l_val + r_th] = 0.0f;
+            }
+        }
+        return;
+    }
+    __syncthreads();
+
+    // Compute column sums of the normalized matrix
+    for (int j = threadIdx.x; j < k; j += blockDim.x) {
+        float sum_col = 0.0f;
+        for (int i = 0; i < k; i++) {
+            sum_col += Y[i * k + j];
+        }
+        col_sum[j] = sum_col;
+    }
+    __syncthreads();
+
+    // Compute conditional probabilities
+    for (int i = threadIdx.x; i < k; i += blockDim.x) {
+        float row_sum = 0.0f;
+        for (int j = 0; j < k; j++) {
+            row_sum += Y[i * k + j];
+        }
+
+        for (int j = 0; j < k; j++) {
+            float cond = 0.0f;
+            if (row_sum != 0.0f) {
+                cond = Y[i * k + j] / row_sum;
+            }
+
+            float final_val = 0.0f;
+            if (col_sum[j] != 0.0f) {
+                final_val = cond / col_sum[j];
+            }
+
+            // Write to output with (row, column, threshold) ordering
+            out[i * (k * l_val) + j * l_val + r_th] = final_val;
+        }
+    }
+    __syncthreads();
 }
 """
-occur_count_kernel2 = cp.RawKernel(kernel_code2, "occur_reduction_kernel")
+occur_reduction_kernel_global = cp.RawKernel(
+    occur_reduction_kernel_code_global, "occur_reduction_kernel_global"
+)