`raise=true` cannot lower an ordinary pivoted Gauss-Jordan solve

[masonlee277]

Summary

An ordinary 3×3 Gauss-Jordan solve with partial pivoting works on the KernelAbstractions CPU backend, but Reactant.@compile raise=true fails while lowering the pivot search. The reproducer uses normal loops over one augmented matrix: there are no generated functions, fully unrolled expressions, static arrays, or explicit matrix inverse.

Reproduced on Reactant 0.2.271 (2026-07-10).

Environment

Julia 1.12.5 · Reactant 0.2.271 · Enzyme 0.13.173 · CUDA 6.2.0 · KernelAbstractions 0.9.42 · CPU backend (CUDA loaded for ReactantCUDAExt)

Reproducer

using Reactant
using CUDA
using KernelAbstractions
import KernelAbstractions as KA

Reactant.set_default_backend("cpu")
@assert Reactant.is_extension_loaded(Val(:CUDA))

@kernel function gauss_jordan_solve!(Ab)
    @inbounds for k in 1:3
        pivot_row = k
        pivot_abs = abs(Ab[k, k])
        for row in (k + 1):3
            candidate = abs(Ab[row, k])
            if candidate > pivot_abs
                pivot_abs = candidate
                pivot_row = row
            end
        end

        if pivot_row != k
            for column in 1:4
                value = Ab[k, column]
                Ab[k, column] = Ab[pivot_row, column]
                Ab[pivot_row, column] = value
            end
        end

        pivot = Ab[k, k]
        for column in k:4
            Ab[k, column] /= pivot
        end
        for row in 1:3
            if row != k
                factor = Ab[row, k]
                for column in k:4
                    Ab[row, column] -= factor * Ab[k, column]
                end
            end
        end
    end
end

function solve!(Ab)
    gauss_jordan_solve!(KA.get_backend(Ab))(Ab; ndrange = 1)
    return nothing
end

Ab0 = [0.0 2.0 1.0 1.0; 1.0 1.0 0.0 2.0; 2.0 1.0 1.0 4.0]
if ARGS == ["--eager"]
    Ab = copy(Ab0)
    gauss_jordan_solve!(KA.CPU())(Ab; ndrange = 1)
    @assert Ab ≈ [1.0 0.0 0.0 5 / 3; 0.0 1.0 0.0 1 / 3; 0.0 0.0 1.0 1 / 3]
else
    Ab = Reactant.to_rarray(Ab0)
    compiled = Reactant.@compile raise=true raise_first=true sync=true solve!(Ab)
    compiled(Ab)
end

The eager control passes:

julia gauss_jordan_solve_raise.jl --eager

Raised compilation fails:

julia gauss_jordan_solve_raise.jl

Expected behavior

The kernel should compile and transform the augmented matrix [A b] into [I x], with x = [5/3, 1/3, 1/3].

Actual behavior

ERROR: LoadError: CompilationError: MLIR pass pipeline "all" failed
error: cannot raise op to stablehlo
%3:3 = "scf.for"(...) ({
  ...
  %candidate = "memref.load"(...)
  %larger = "arith.cmpf"(...)
  %pivot_row = "arith.select"(%larger, %row, %previous_row)
  "scf.yield"(..., %pivot_row, %pivot_row)
})

The first reported source location is the runtime pivot search ending at line 19 of the reproducer.

Why this matters

This is the compact solve-oriented form we would like to use for small pivoted systems inside raised kernels. Expanding the algorithm into a fully unrolled scalar expression is not a practical substitute because it makes compiler IR grow with the matrix dimension.

Related issues

Related control-flow raising gaps are EnzymeAD/Enzyme-JAX #442 and #1672.