[RFC] Refactoring Scanpy's Numba Backend: C-Style Kernel Layer & Biology-Optimized Sparse Matrices

[krkawzq]

What kind of feature would you like to request?

Additional function parameters / changed functionality / changed defaults?

Please describe your wishes

1. Motivation: Current State of Scanpy's Numba Backend

This RFC proposes a gradual refactoring of Scanpy's Numba backend to replace the current ad-hoc, poorly optimized JIT implementations with a structured, high-performance kernel layer.

The Problem: Chaotic Numba Usage

The current Numba usage in Scanpy is disorganized and inefficient. A prime example is scanpy/preprocessing/_highly_variable_genes.py:

# Current Scanpy pattern - problematic
@njit(cache=True, parallel=False)  # ← Parallelism DISABLED
def _sum_and_sum_squares_clipped(indices, data, n_cols, clip_val, nnz):
    squared_batch_counts_sum = np.zeros(n_cols, dtype=np.float64)
    batch_counts_sum = np.zeros(n_cols, dtype=np.float64)
    for i in range(nnz):
        val = data[i]
        if val > clip_val[indices[i]]:
            val = clip_val[indices[i]]
        squared_batch_counts_sum[indices[i]] += val ** 2  # ← Race condition!
        batch_counts_sum[indices[i]] += val               # ← Race condition!
    return squared_batch_counts_sum, batch_counts_sum

Critical Issues:

Problem	Description
Race Conditions	Parallelism disabled (parallel=False) because race conditions on shared arrays were not properly handled via reduction or atomics
No Optimization Hints	Missing assume(), vectorize(), likely() → LLVM cannot optimize aggressively
Pythonic Style	High-level constructs prevent SIMD vectorization and loop unrolling
Tight Coupling	Business logic (string methods, flavor dispatch) mixed into JIT functions → Object mode fallbacks
Unverified Compilation	No CI checks for optimal parallel patterns or Python C-API callbacks

This pattern repeats across Scanpy's codebase, leaving 1-2 orders of magnitude of performance on the table.

---

2. Proposal: A Decoupled Kernel Architecture

2.1 The Kernel Layer

We propose establishing a dedicated kernel layer containing pure computational kernels written in C-style Python. All kernels must satisfy the following requirements (enforced via CI):

Requirement 1: Optimal Parallelization (CI-Verified)

Use numba.parallel_diagnostics() to verify kernels achieve optimal parallel patterns:

# CI check example
from numba import njit, prange

@njit(parallel=True)
def kernel(...):
    for i in prange(n):  # Must show as "PARALLEL" in diagnostics
        ...

# Verify in CI:
# kernel.parallel_diagnostics(level=4) should show no "NOT PARALLEL" warnings

Requirement 2: No Python Callbacks (CI-Verified)

Use inspect_llvm() to verify zero Python C-API calls in compiled code:

# CI check: No calls to scipy.loess, statsmodels, or any interpreted code
llvm_ir = kernel.inspect_llvm()
assert 'NRT_' not in llvm_ir  # No Numba runtime calls for Python objects
assert 'PyObject' not in llvm_ir  # No Python object manipulation

Requirement 3: LLVM Optimization Hints

Use compiler hints to enable optimizations that exceed hand-written C/C++:

from biosparse.optim import parallel_jit, assume, likely, vectorize, unroll

@parallel_jit(boundscheck=False)
def optimized_kernel(csr, n_targets):
    n_rows = csr.nrows
    
    # Tell LLVM what we know → eliminates defensive code
    assume(n_rows > 0)
    assume(n_targets > 0)
    
    for row in prange(n_rows):
        values, indices = csr.row_to_numpy(row)
        
        total = 0.0
        vectorize(8)  # SIMD hint
        unroll(4)     # Loop unrolling hint
        for j in range(len(values)):
            if likely(values[j] > 0):  # Branch prediction
                total += values[j]

Reference implementation: biosparse/optim

Requirement 4: Strict Decoupling

Kernels must be pure computation with no business logic:

# ✅ GOOD: Pure kernel - accepts only primitive data structures
@parallel_jit
def mwu_core(csr, group_ids, group_counts, n_targets, out_U, out_tie):
    ...

# ❌ BAD: Mixed with business logic
@njit
def compute_hvg(adata, flavor='seurat', n_top_genes=2000):  # String params!
    if flavor == 'seurat':  # Business logic in kernel!
        ...

Kernel Inputs: Only NumPy arrays, SciPy CSR/CSC components, or biosparse structures.
No: Strings, AnnData objects, flavor switches, parameter validation.

2.2 The Frontend Layer

Refactor scanpy.tools and scanpy.preprocessing as pure dispatchers:

┌─────────────────────────────────────────────────────────────────┐
│  FRONTEND: scanpy.preprocessing / scanpy.tools                  │
│  ─────────────────────────────────────────────────────────────  │
│  • Parse arguments & validate parameters                        │
│  • Handle flavor selection (seurat, cell_ranger, etc.)          │
│  • AnnData I/O (read adata.X, write adata.var)                  │
│  • Convert data structures for kernel consumption               │
│  • NO mathematical computation                                   │
└───────────────────────────┬─────────────────────────────────────┘
                            │ dispatch
                            ▼
┌─────────────────────────────────────────────────────────────────┐
│  BACKEND: kernel layer (new module)                              │
│  ─────────────────────────────────────────────────────────────  │
│  • Pure C-style Numba JIT kernels                               │
│  • Verified parallel patterns (CI)                               │
│  • Verified no Python callbacks (CI)                             │
│  • LLVM optimization hints (assume, vectorize, etc.)            │
│  • Only primitive data structures                                │
└─────────────────────────────────────────────────────────────────┘

---

3. Sparse Matrix: Biology-First Data Structure

The Problem with SciPy Sparse

SciPy's CSR/CSC matrices are designed for general linear algebra, not biological data analysis. Key limitations:

Operation	SciPy Behavior	Biology Need
Slicing X[1000:2000, :]	Copies data	Zero-cost view
Stacking vstack([A, B])	Allocates new arrays	Efficient concatenation
Numba integration	Requires unpacking to arrays	Native JIT support
Memory layout	General-purpose	Optimized for cell×gene access

Recommendation: biosparse

biosparse provides sparse matrices designed for biological data:

from biosparse import CSRF64

# Zero-copy from scipy
csr = CSRF64.from_scipy(scipy_mat, copy=False)

# Zero-cost slicing (views, not copies)
subset = csr[1000:2000, :]       # No data copy!
genes_subset = csr[:, gene_mask] # No data copy!

# Efficient stacking for dataset merging
merged = CSRF64.vstack([dataset1, dataset2, dataset3])

# Native Numba integration
@parallel_jit
def my_kernel(csr: CSR):
    for row in prange(csr.nrows):
        values, indices = csr.row_to_numpy(row)  # Direct access
        ...

Designed for million-scale cell datasets where SciPy's copy-on-slice becomes a critical bottleneck.

---

4. Case Study: biosparse

I have implemented biosparse as a proof-of-concept demonstrating this architecture. It is a case study project implementing selected kernels to validate the proposed approach.

4.1 Implemented Kernels

Currently implemented (for case study purposes):

Kernel	Description	Location
HVG	Seurat, Seurat V3, Cell Ranger, Pearson residuals	kernel/hvg.py
MWU	Mann-Whitney U test	kernel/mwu.py
t-test	Welch's and Student's t-test	kernel/ttest.py

All kernels: biosparse/kernel

4.2 Performance Results

biosparse achieves 20-100x speedup over Scanpy's current implementations:

Kernel	Speedup vs Scanpy
HVG (Seurat)	30x
HVG (Seurat V3)	30x
HVG (Pearson)	80x
Mann-Whitney U	25x

Benchmarks: biosparse/benchmarks

4.3 Implementation Style

All kernels are written in pure Python but follow C-style conventions:

@parallel_jit(cache=True, boundscheck=False)
def _mwu_core(
    csr: CSR,
    group_ids: np.ndarray,
    group_counts: np.ndarray,
    n_targets: int,
    out_U1: np.ndarray,
    out_tie: np.ndarray,
    out_sum_ref: np.ndarray,
    out_sum_tar: np.ndarray,
    # Thread-local pre-allocated buffers (indexed by get_thread_id())
    tl_buf_ref: np.ndarray,      # (n_threads, max_nnz + 1)
    tl_buf_tar: np.ndarray,      # (n_threads, n_targets, max_nnz + 1)
    tl_n_tar_nz: np.ndarray,     # (n_threads, n_targets)
    tl_sum_tar: np.ndarray,      # (n_threads, n_targets)
) -> None:
    """Core MWU - C-style optimized kernel with thread-local buffers.
    
    Thread-local buffers eliminate heap allocation overhead in prange loop.
    Buffers are indexed by get_thread_id() for zero-contention access.
    """
    n_rows = csr.nrows
    n_ref = group_counts[0]
    
    # Precompute constants (avoid per-iteration computation)
    n_ref_f = float(n_ref)
    half_n1_n1p1 = 0.5 * n_ref_f * (n_ref_f + 1.0)
    
    # Compiler hints
    assume(n_rows > 0)
    assume(n_targets > 0)
    assume(n_ref > 0)
    
    for row in prange(n_rows):
        # Get thread-local buffer via thread ID (ZERO heap allocation!)
        tid = get_thread_id()
        buf_ref = tl_buf_ref[tid]
        buf_tar = tl_buf_tar[tid]
        n_tar_nz = tl_n_tar_nz[tid]
        sum_tar = tl_sum_tar[tid]
        
        values, col_indices = csr.row_to_numpy(row)
        nnz = len(values)
        
        n_ref_nz = 0
        sum_ref = 0.0
        
        # Reset thread-local counters (no allocation, just zero-init)
        vectorize(4)
        unroll(4)
        for t in range(n_targets):
            n_tar_nz[t] = 0
            sum_tar[t] = 0.0
        
        # Single-pass scan with branch prediction
        for j in range(nnz):
            col = col_indices[j]
            val = float(values[j])  # Explicit cast
            g = group_ids[col]
            
            if g == 0:
                buf_ref[n_ref_nz] = val
                sum_ref += val
                n_ref_nz += 1
            elif likely(g > 0):
                if likely(g <= n_targets):
                    t = g - 1
                    buf_tar[t, n_tar_nz[t]] = val
                    sum_tar[t] += val
                    n_tar_nz[t] += 1
        
        # ... (sorting and rank computation)

C-Style Optimization Techniques Used:

• Thread-local buffers - pre-allocate once, index by get_thread_id() to eliminate heap allocation in prange
• All constants inlined / precomputed
• assume() for bounds elimination
• likely()/unlikely() for branch prediction
• vectorize()/unroll() for SIMD and ILP
• np.empty() + manual init (faster than np.zeros in prange)
• Explicit type casts to avoid implicit conversions
• Insertion sort for small arrays, numpy sort for large
• Binary search instead of linear scan
• Single-pass algorithms where possible

---

5. Technical Challenge: Dynamic Loop Hints

The Problem

biosparse's optim module is currently a lightweight toolkit. Numba's support for loop hints is limited, particularly around dynamic vectorization width dispatch.

Current approach uses hardcoded hints:

vectorize(8)  # Hardcoded 8-wide SIMD
for i in range(n):
    ...

This can confuse LLVM's middle-end during type-based dispatch. For example, an f64 kernel entering a loop marked vectorize(8) assumes 512-bit vectors (AVX-512), but:

• f32 data could use 16 lanes at the same vector width
• Not all CPUs support AVX-512
• LLVM's cost model might prefer different widths

Seeking Input

Since biosparse intentionally avoids C/C++ extensions, we need a pure Numba/Python solution for:

1. Compile-time type inspection to determine optimal vector width
2. Runtime CPU feature detection for SIMD capability dispatch
3. Better integration with LLVM's auto-vectorization cost model

We welcome suggestions from maintainers on solving this within the Numba ecosystem.

Optimization toolkit: biosparse/optim

---

6. Conclusion

biosparse demonstrates that Scanpy's current Numba usage has 1-2 orders of magnitude optimization headroom. The existing codebase suffers from:

• Race conditions handled by disabling parallelism
• No compiler optimization hints
• Tight coupling between business logic and computation
• Suboptimal data structures for biological workloads

Proposed Actions

1. Establish a kernel layer with CI-enforced quality standards
2. Refactor frontend modules to pure dispatchers
3. Adopt biology-optimized sparse matrices (e.g., biosparse)
4. Incrementally migrate kernels to C-style implementations

Resources

Resource	Link
biosparse repository	https://github.com/krkawzq/biosparse
Kernel implementations	biosparse/kernel
Optimization toolkit	biosparse/optim
Benchmarks	biosparse/benchmarks

---

biosparse is a case study project I developed to demonstrate this architecture. I'm happy to discuss implementation details, contribute to refactoring efforts, or provide additional benchmarks.

[ilan-gold]

Hi @krkawzq to help us get to know you and your process a bit maybe we should start with something small (after all, what you are proposing is incremental).

@njit(cache=True, parallel=False) # ← Parallelism DISABLED

It says in the codebase this is for accuracy purposes but you seem to say it's "because race conditions on shared arrays were not properly handled via reduction or atomics." Is there any way to resolve this problem by altering the current implementation? Could you submit a small PR to fix this or is it impossible?

Could you also shortly highlight what advantage your CSR construct offers over scipy sparse matrices? At the end of the day, I can't imagine what it does differently given that we support sparse matrices in numba as types natively now in fau and thus in scanpy: see

scanpy/src/scanpy/experimental/pp/_highly_variable_genes.py

Lines 34 to 36 in d5901ab

	@njit
	def _calculate_res_sparse(
	mat: CSBase,

for an example, although I'm sure it's use could be expanded. For the purpose of pure speed (which I'm trying to focus on in this reply), I don't see much in your comparison that would distinguish the two maybe except memory layout given out support for numba.

Thanks for the issue!

[krkawzq]

@ilan-gold
Hi,

Thank you for the thoughtful feedback and for taking the time to engage with this proposal! I really appreciate your pragmatic approach of starting small. Let me address your questions directly:

1. Regarding the Parallelization Issue: This Is a Race Condition, Not a Precision Trade-off

The comment in the codebase stating # parallel=False needed for accuracy is misleading. This is not a floating-point precision issue—it's a classic race condition that produces incorrect results, not merely imprecise ones.

Why This Is a Race Condition

The indices array contains column indices from a sparse matrix, meaning the same index can appear multiple times. When parallel=True is enabled, multiple threads may simultaneously execute:

squared_batch_counts_sum[idx] += element**2

The += operator is a non-atomic read-modify-write operation. When two threads access the same idx concurrently, one thread's write can be completely lost—this isn't imprecision, it's a data race that silently corrupts results.

Parallelism Cannot Cause Precision Loss

To be absolutely clear: parallelism itself cannot introduce precision errors. The only floating-point concern with parallel execution comes from fastmath, which may reorder operations (affecting reproducibility due to non-associativity of floating-point addition). However, that's a separate issue and doesn't justify disabling parallelism.

The results produced by the current implementation with parallel=False aren't "more precise"—they're simply the only correct results because the race condition is avoided by serialization.

Suggested Fix

Yes, this can be fixed while maintaining parallelism. Here's a corrected implementation using Numba's parallel reduction pattern:

@numba.njit(cache=True, parallel=True)
def _sum_and_sum_squares_clipped(
    indices: NDArray[np.integer],
    data: NDArray[np.floating],
    *,
    n_cols: int,
    clip_val: NDArray[np.float64],
    nnz: int,
) -> tuple[NDArray[np.float64], NDArray[np.float64]]:
    """
    Parallel implementation using thread-local buffers to avoid race conditions.
    
    Previous implementation used parallel=False due to race condition on shared arrays.
    This version uses explicit thread-local reduction to restore both correctness 
    and parallelism.
    """
    # Thread-local accumulators for parallel reduction
    n_threads = numba.get_num_threads()
    squared_local = np.zeros((n_threads, n_cols), dtype=np.float64)
    sum_local = np.zeros((n_threads, n_cols), dtype=np.float64)
    
    # Parallel accumulation into thread-local buffers (no race condition)
    for i in numba.prange(nnz):
        tid = numba.get_thread_id()
        idx = indices[i]
        element = min(np.float64(data[i]), clip_val[idx])
        squared_local[tid, idx] += element**2
        sum_local[tid, idx] += element
    
    # Reduction phase: combine thread-local results
    squared_batch_counts_sum = np.zeros(n_cols, dtype=np.float64)
    batch_counts_sum = np.zeros(n_cols, dtype=np.float64)
    
    for t in range(n_threads):
        for j in range(n_cols):
            squared_batch_counts_sum[j] += squared_local[t, j]
            batch_counts_sum[j] += sum_local[t, j]
    
    return squared_batch_counts_sum, batch_counts_sum

How this fixes the issue:

• Each thread accumulates into its own buffer (squared_local[tid, :])
• No two threads ever write to the same memory location during the parallel loop
• A single-threaded reduction phase combines the results afterward

I'd be happy to submit this as a small PR with unit tests to verify:

1. Results match the original implementation exactly
2. Parallelism is properly enabled
3. Performance improvement is measurable

---

2. Regarding biosparse vs. FAU/SciPy Sparse Matrices

You're absolutely right to question this—I want to clarify that biosparse and the kernel refactoring are independent proposals. The kernel layer refactoring is the core proposal; biosparse was presented as a supporting case study, but it's not a prerequisite.

That said, let me briefly explain the memory layout difference you mentioned, since it does have practical implications for certain bioinformatics workflows.

The Core Difference: Contiguous vs. Non-Contiguous Memory

You correctly identified that the key distinction is memory layout. Let me illustrate what this means in practice:

SciPy CSR (contiguous memory):

indptr:  [0, 2, 4, 6]
data:    [█████████████████████████████]  # single contiguous allocation
indices: [█████████████████████████████]  # single contiguous allocation

All non-zero values across all rows live in one contiguous data array. Row boundaries are encoded implicitly via indptr.

BioSparse (non-contiguous, per-row memory):

values:  [Span₀] [Span₁] [Span₂] ...
            ↓       ↓       ↓
         [███]   [███]   [███]   # each row has its own memory region

Each row is stored as an independent memory region (Span). The matrix is essentially a vector of row pointers.

Why This Matters: Zero-Copy Slicing and Stacking

The practical consequence shows up in operations common to bioinformatics:

Operation 1: Row Subsetting (e.g., selecting 1000 cells from a 50k cell dataset)

# SciPy - O(nnz) data copy required
subset = csr_matrix[selected_cells, :]  # Must rebuild indptr, copy data/indices

# BioSparse - O(k) pointer clones, no data copy
subset = csr[selected_cells, :]  # Just clone row references (Arc::clone)

Operation 2: Dataset Merging (e.g., combining 3 batches into one dataset)

# SciPy - O(total_nnz) allocation + copy
merged = scipy.sparse.vstack([batch1, batch2, batch3])  # Allocates new arrays

# BioSparse - O(n_rows) pointer moves, no data copy
merged = CSRF64.vstack([batch1, batch2, batch3])  # Just concatenate row pointers

Cost comparison:

Operation	SciPy	BioSparse
Row slice (k rows, n nnz per row)	O(k × n) copy	O(k) Arc clones
vstack (m matrices, N total nnz)	O(N) copy	O(total_rows) moves

The Fragmentation Problem and Span-Storage Architecture

You might wonder: "Wouldn't per-row allocation cause severe memory fragmentation?"

Yes, it would—which is why biosparse uses a Span-Storage joint architecture:

┌─────────────────────────────────────────────────────────┐
│             Storage (Arc-counted memory block)          │
│  ┌─────────────────────────────────────────────────┐    │
│  │  [Row0 data | Row1 data | Row2 data | ...]      │    │
│  └─────────────────────────────────────────────────┘    │
└─────────────────────────────────────────────────────────┘
         ↑           ↑           ↑
      Span₀       Span₁       Span₂  (lightweight views)

Multiple Span objects (lightweight row views) share a single contiguous Storage block via Arc reference counting. This achieves:

• Zero fragmentation: Bulk allocation of contiguous memory
• Logical separation: Each row is an independent object
• Shared ownership: Views share memory without copying (reference-counted)

When you slice or stack matrices, you're just cloning/moving these lightweight Span references—the underlying data stays put.

Important Caveat: View Semantics

Because slicing produces views (shared ownership), not copies:

subset = csr[1000:2000, :]  # subset shares memory with csr
# They point to the same underlying Storage

If you need an independent copy (e.g., to modify without affecting the original), you must explicitly deep copy:

independent = subset.deep_copy()  # Allocates new memory

This is different from SciPy, where slicing usually produces copies by default.

---

My Honest Assessment

Having explained all this, I want to be clear about my priorities:

biosparse addresses the smallest bottleneck among the issues I raised. Yes, it's faster for row slicing and stacking—but these operations aren't the main performance bottleneck in typical Scanpy workflows.

The more critical optimizations are:

1. Fixing race conditions (like the one we discussed above)
2. Proper parallel kernel design with LLVM optimization hints
3. Memory-mapped operators for out-of-core computation
4. Efficient disk-to-memory loading via masks

Regarding FAU: You mentioned that FAU provides native Numba support for sparse matrices via CSBase. I agree this is excellent progress. biosparse essentially provides a similar Numba integration, but with the non-contiguous layout. The question is whether the zero-copy slicing benefits justify:

• Introducing a new dependency
• Learning a different API
• Maintaining compatibility with SciPy ecosystem

Given SciPy's ecosystem advantages (existing tools, familiarity, stability), I'm not advocating for wholesale replacement. biosparse was built as a case study to validate the kernel optimization approach—it proved that C-style kernels can achieve 20-100x speedups regardless of the underlying sparse format.

If the community prefers sticking with SciPy/FAU and focusing purely on computational kernel optimization (the main proposal), I'm fully supportive of that direction. The kernel refactoring doesn't depend on biosparse.

[ilan-gold]

The more critical optimizations are:

This is a good prioritization list and will also help a bit more clearly one-to-one compare to assess biosparse's sparse format.

Yes, this can be fixed while maintaining parallelism. Here's a corrected implementation using Numba's parallel reduction pattern:

So I think this would be a great first contribution. If we see a nice perf boost, we can start looking into improving all of our numba kernels - I don't see any reason why the optimizations you have in biosparse (like vectorize) can't be applied to our kernels!

Thanks a million, great issue!