First Numba CPU Ports: Mantel and PERMANOVA

Merged Pull Requests

Pull Request in Active Review

Numba 0.65.1 parfor cycle bug

The natural Numba code for PERMANOVA's s_W crashes at compile time with AssertionError: unexpected cycle in lookup() from inside the parfor reduction-analysis pass. The crash needs two features together: a cross-iteration scalar reduction in the outer prange, and an if conditional inside the inner loop that reads from a parameter array. Either feature alone compiles fine.

@njit(parallel=True)
def natural_sW(mat, group_sizes, grouping):
    n = mat.shape[0]
    s_W = 0.0
    for row in prange(n - 1):
        group = grouping[row]
        local = 0.0
        for col in range(row + 1, n):
            if grouping[col] == group:        # conditional read of a parameter
                val = mat[row, col]
                local += val * val
        s_W += local / group_sizes[group]      # cross-iteration scalar reduction
    return s_W

The workaround replaces the scalar reduction with per-iteration writes to a partials array, summed once after the prange exits. The conditional stays in place. Beyond sidestepping the bug, this pattern has a useful property for the future GPU work: parallel=True reductions are a CPU-only Numba feature, but the per-iteration-slot pattern maps directly to @cuda.jit.

Engineering lesson: tests should validate against trusted references

One review point that came up on PR #2470 was about test design. The initial test class included a Python reimplementation of the s_W computation as the "expected value" for the paired _cy and _numba tests. That carries a subtle bug-masking risk: if the test reimplementation and the function under test share the same misunderstanding of the formula, both pass together while being wrong together.

The fix looks obvious in hindsight. Compute the expected value once using the trusted reference implementation (here, the pre-existing Cython), hardcode the resulting number, and assert against it. The test no longer carries an algorithmic implementation. It carries a single value from a trusted source.

Peer collaboration on PCoA

Another GSoC contributor on the same scikit-bio project is working on Numba CPU and GPU implementations of center_distance_matrix for PCoA. The first week of collaboration covered how to structure the work as separate CPU and GPU PRs, how to slot the new tests into the existing Numba CI workflow, and how the Mantel and PERMANOVA architecture can serve as a template for that branch.

What's queued next

Three pieces of work are queued behind PR #2470:

1. Interface cleanup PR. Make Numba opt-in rather than auto-selected on installation, and rename the _numba function suffix to _nb to match the existing _cy convention. This will target Mantel and PERMANOVA at the same time.
2. Block-permutation optimization for PERMANOVA. Process multiple permutations per matrix scan, with cache tiling and a parallelization axis flipped from matrix rows to permutation blocks. The C++ reference for this is in the scikit-bio-binaries repository.
3. GPU @cuda.jit implementation. Builds on the optimized CPU code. The partials pattern from the parfor workaround already aligns with the structure needed for this transition.