Use MPI jail escape for precompilation

Replace SnoopCompile static precompile.jl with PrecompileTools + MPI
environment cleanup. The precompilation subprocess clears PMI_*, PMIX_*,
OMPI_*, and MPI_* environment variables before MPI.Init(), allowing it
to initialize as a fresh single-rank process even when the parent was
launched under mpiexec.

Author: Sébastien Loisel

Project.toml

@@ -8,18 +8,19 @@ Blake3Hash = "8f478455-a32d-4928-b0e4-72b19a7d5574"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 MPI = "da04e1cc-30fd-572f-bb4f-1f8673147195"
 MUMPS = "55d2b088-9f4e-11e9-26c0-150b02ea6a46"
+PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 
 [compat]
 Blake3Hash = "0.3"
 MPI = "0.20"
 MUMPS = "1.5"
+PrecompileTools = "1"
 julia = "1.10"
 
 [extras]
 BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
-SnoopCompile = "aa65fe97-06da-5843-b5b1-d5d13cad87d2"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]

scripts/generate_precompile.jl

@@ -662,8 +662,135 @@ end
 # Precompilation Workload
 # ============================================================================
 
-# Precompile directives generated by SnoopCompile
-# Regenerate with: mpiexec -n 1 julia --project=. scripts/generate_precompile.jl
-include("precompile.jl")
+using PrecompileTools
+
+@setup_workload begin
+    # Small test data for precompilation (runs with single MPI rank)
+    n = 8
+
+    # Sparse matrix (tridiagonal) - Float64
+    I_sp = Int[]; J_sp = Int[]; V_sp = Float64[]
+    for i in 1:n
+        push!(I_sp, i); push!(J_sp, i); push!(V_sp, 4.0)
+        if i > 1
+            push!(I_sp, i); push!(J_sp, i-1); push!(V_sp, -1.0)
+            push!(I_sp, i-1); push!(J_sp, i); push!(V_sp, -1.0)
+        end
+    end
+    A_sparse_f64 = sparse(I_sp, J_sp, V_sp, n, n)
+
+    # Sparse matrix - ComplexF64
+    A_sparse_c64 = sparse(I_sp, J_sp, ComplexF64.(V_sp), n, n)
+
+    # Dense matrix - Float64
+    A_dense_f64 = Float64[i == j ? 4.0 : (abs(i-j) == 1 ? -1.0 : 0.0) for i in 1:n, j in 1:n]
+
+    # Dense matrix - ComplexF64
+    A_dense_c64 = ComplexF64.(A_dense_f64)
+
+    # Vectors
+    v_f64 = ones(Float64, n)
+    v_c64 = ones(ComplexF64, n)
+
+    # Identity for SPD construction
+    I_sparse = sparse(1.0 * LinearAlgebra.I, n, n)
+
+    @compile_workload begin
+        # === MPI Jail Escape ===
+        # When precompiling under mpiexec, the subprocess inherits MPI environment
+        # variables but isn't part of the MPI job. Clean them to allow MPI.Init()
+        # to succeed as a fresh single-rank process.
+        for k in collect(keys(ENV))
+            if startswith(k, "PMI") || startswith(k, "PMIX") || startswith(k, "OMPI_") || startswith(k, "MPI_")
+                delete!(ENV, k)
+            end
+        end
+
+        MPI.Init()
+
+        # === VectorMPI operations (Float64) ===
+        v = VectorMPI(v_f64)
+        w = VectorMPI(2.0 .* v_f64)
+        _ = v + w
+        _ = v - w
+        _ = 2.0 * v
+        _ = v * 2.0
+        _ = norm(v)
+        _ = dot(v, w)
+        _ = conj(v)
+        _ = length(v)
+        _ = size(v)
+
+        # VectorMPI (ComplexF64)
+        vc = VectorMPI(v_c64)
+        _ = conj(vc)
+        _ = norm(vc)
+
+        # === SparseMatrixMPI operations (Float64) ===
+        A = SparseMatrixMPI{Float64}(A_sparse_f64)
+        B = SparseMatrixMPI{Float64}(A_sparse_f64)
+        _ = A + B
+        _ = A - B
+        _ = 2.0 * A
+        _ = A * v
+        _ = A * B
+        _ = transpose(A)
+        At = SparseMatrixMPI(transpose(A))
+        _ = size(A)
+        _ = nnz(A)
+        _ = norm(A)
+
+        # SparseMatrixMPI (ComplexF64)
+        Ac = SparseMatrixMPI{ComplexF64}(A_sparse_c64)
+        _ = Ac * vc
+
+        # === MatrixMPI operations (Float64) ===
+        D = MatrixMPI(A_dense_f64)
+        _ = 2.0 * D
+        _ = D * v
+        _ = transpose(D)
+        Dt = copy(transpose(D))  # Materialize dense transpose
+        _ = size(D)
+        _ = norm(D)
+
+        # MatrixMPI (ComplexF64)
+        Dc = MatrixMPI(A_dense_c64)
+        _ = Dc * vc
+
+        # === Mixed operations ===
+        _ = A * D  # Sparse * Dense
+
+        # === Indexing ===
+        _ = v[1]
+        _ = A[1, 1]
+        _ = D[1, 1]
+
+        # === Factorization (MUMPS) ===
+        # Make symmetric positive definite: A + A^T + 10I
+        At_mat = SparseMatrixMPI(transpose(A))
+        I_dist = SparseMatrixMPI{Float64}(I_sparse)
+        A_spd = A + At_mat + I_dist * 10.0
+        F = LinearAlgebra.ldlt(A_spd)
+        x = F \ v
+        finalize!(F)
+
+        # LU factorization
+        F_lu = LinearAlgebra.lu(A)
+        x = F_lu \ v
+        finalize!(F_lu)
+
+        # === Block operations ===
+        _ = cat(v, w; dims=1)
+        _ = blockdiag(A, B)
+
+        # === Conversions ===
+        _ = Vector(v)
+        _ = Matrix(D)
+        _ = SparseMatrixCSC(A)
+
+        # Clear caches
+        clear_plan_cache!()
+    end
+end
 
 end # module LinearAlgebraMPI