Unify CPU/GPU code paths with rowptr_target/colval_target

Replace conditional GPU array checks with always-populated target arrays:
- For CPU: rowptr_target === rowptr, colval_target === colval (same objects)
- For GPU: rowptr_target/colval_target are GPU copies

This eliminates runtime conditionals in SpMV and other operations,
providing a single unified code path for both CPU and GPU backends.

Changes:
- Rename rowptr_gpu/colval_gpu to rowptr_target/colval_target
- Add _to_target_backend() helper function
- Update all SparseMatrixMPI constructor sites (~20 locations)
- Update SpMV to use target arrays directly without conditionals
- Update Metal extension mtl()/cpu() functions

Author: Sébastien Loisel

ext/LinearAlgebraMPIMetalExt.jl

@@ -41,11 +41,14 @@ const MtlSparseMatrixMPI{T,Ti} = LinearAlgebraMPI.SparseMatrixMPI{T,Ti,MtlVector
     mtl(A::LinearAlgebraMPI.SparseMatrixMPI)
 
 Convert a CPU SparseMatrixMPI to a Metal GPU SparseMatrixMPI.
-Only the nonzero values (`nzval`) are moved to GPU. Structural arrays
-(`rowptr`, `colval`, partitions) remain on CPU.
+The `nzval` and target structure arrays are moved to GPU.
+The CPU structure arrays (`rowptr`, `colval`, partitions) remain on CPU for MPI.
 """
 function LinearAlgebraMPI.mtl(A::LinearAlgebraMPI.SparseMatrixMPI{T,Ti,Vector{T}}) where {T,Ti}
     nzval_gpu = MtlVector(A.nzval)
+    # Convert structure arrays to GPU (used by unified SpMV kernel)
+    rowptr_target = MtlVector(A.rowptr)
+    colval_target = MtlVector(A.colval)
     return LinearAlgebraMPI.SparseMatrixMPI{T,Ti,MtlVector{T}}(
         A.structural_hash,
         A.row_partition,
@@ -57,7 +60,9 @@ function LinearAlgebraMPI.mtl(A::LinearAlgebraMPI.SparseMatrixMPI{T,Ti,Vector{T}
         A.nrows_local,
         A.ncols_compressed,
         nothing,  # Invalidate cached_transpose (would need to convert too)
-        A.cached_symmetric
+        A.cached_symmetric,
+        rowptr_target,
+        colval_target
     )
 end
 
@@ -68,6 +73,7 @@ Convert a Metal GPU SparseMatrixMPI to a CPU SparseMatrixMPI.
 """
 function LinearAlgebraMPI.cpu(A::LinearAlgebraMPI.SparseMatrixMPI{T,Ti,<:MtlVector}) where {T,Ti}
     nzval_cpu = Array(A.nzval)
+    # For CPU, rowptr_target and colval_target are the same as rowptr and colval
     return LinearAlgebraMPI.SparseMatrixMPI{T,Ti,Vector{T}}(
         A.structural_hash,
         A.row_partition,
@@ -79,7 +85,9 @@ function LinearAlgebraMPI.cpu(A::LinearAlgebraMPI.SparseMatrixMPI{T,Ti,<:MtlVect
         A.nrows_local,
         A.ncols_compressed,
         nothing,  # Invalidate cached_transpose
-        A.cached_symmetric
+        A.cached_symmetric,
+        A.rowptr,  # rowptr_target (same as rowptr for CPU)
+        A.colval   # colval_target (same as colval for CPU)
     )
 end
 
@@ -140,4 +148,24 @@ function LinearAlgebraMPI._create_output_like(v::LinearAlgebraMPI.VectorMPI{T,<:
     return LinearAlgebraMPI.mtl(v)
 end
 
+# ============================================================================
+# MatrixPlan Index Array Support
+# ============================================================================
+
+"""
+    _index_array_type(::Type{MtlVector{T}}, ::Type{Ti}) where {T,Ti}
+
+Map MtlVector{T} value array type to MtlVector{Ti} index array type.
+Used by MatrixPlan to store symbolic index arrays on GPU.
+"""
+LinearAlgebraMPI._index_array_type(::Type{MtlVector{T}}, ::Type{Ti}) where {T,Ti} = MtlVector{Ti}
+
+"""
+    _to_target_backend(v::Vector{Ti}, ::Type{MtlVector{T}}) where {Ti,T}
+
+Convert a CPU index vector to Metal GPU.
+Used by SparseMatrixMPI constructors to create GPU structure arrays.
+"""
+LinearAlgebraMPI._to_target_backend(v::Vector{Ti}, ::Type{<:MtlVector}) where {Ti} = MtlVector(v)
+
 end # module

src/LinearAlgebraMPI.jl

@@ -109,8 +109,8 @@ function SparseArrays.SparseMatrixCSC(A::SparseMatrixCSR{Tv,Ti}) where {Tv,Ti}
 end
 
 # Cache for memoized MatrixPlans
-# Key: (A_hash, B_hash, T, Ti) - use full 256-bit hashes
-const _plan_cache = Dict{Tuple{Blake3Hash,Blake3Hash,DataType,DataType},Any}()
+# Key: (A_hash, B_hash, T, Ti, AV) - use full 256-bit hashes, includes array type for GPU support
+const _plan_cache = Dict{Tuple{Blake3Hash,Blake3Hash,DataType,DataType,Type},Any}()
 
 # Cache for memoized VectorPlans (for A * x)
 # Key: (A_hash, x_hash, T, AV) - includes array type for GPU support
@@ -141,8 +141,8 @@ end
 const _diag_structure_cache = Dict{Blake3Hash,DiagStructureCache}()
 
 # Cache for memoized AdditionPlans (for A + B and A - B)
-# Key: (A_hash, B_hash, T, Ti) - use full 256-bit hashes
-const _addition_plan_cache = Dict{Tuple{Blake3Hash,Blake3Hash,DataType,DataType},Any}()
+# Key: (A_hash, B_hash, T, Ti, AV) - use full 256-bit hashes, includes array type for GPU support
+const _addition_plan_cache = Dict{Tuple{Blake3Hash,Blake3Hash,DataType,DataType,Type},Any}()
 
 # Cache for memoized IdentityAdditionPlans (for A + λI)
 # Key: (A_hash, T, Ti) - use full 256-bit hash

src/indexing.jl

@@ -1049,9 +1049,10 @@ function Base.getindex(A::SparseMatrixMPI{T,Ti}, row_rng::UnitRange{Int}, col_rn
         # SparseMatrixCSC(ncols, nrows, colptr, rowval, nzval) - transposed storage
         empty_AT = SparseMatrixCSC(new_ncols, my_local_rows, ones(Int, my_local_rows + 1), Int[], T[])
         hash = compute_structural_hash(new_row_partition, Int[], empty_AT, comm)
+        # For CPU, rowptr_target and colval_target are the same as rowptr and colval
         return SparseMatrixMPI{T,Ti,Vector{T}}(hash, new_row_partition, new_col_partition, Int[],
                                    empty_AT.colptr, empty_AT.rowval, empty_AT.nzval,
-                                   my_local_rows, 0, nothing, nothing)
+                                   my_local_rows, 0, nothing, nothing, empty_AT.colptr, empty_AT.rowval)
     end
 
     # Compute new row partition (local computation, no communication)
@@ -1161,9 +1162,10 @@ function Base.getindex(A::SparseMatrixMPI{T,Ti}, row_rng::UnitRange{Int}, col_rn
     hash = compute_structural_hash(new_row_partition, final_col_indices, new_AT.colptr, new_AT.rowval, comm)
 
     # Use explicit CSR arrays for the new struct format
+    # For CPU, rowptr_target and colval_target are the same as rowptr and colval
     return SparseMatrixMPI{T,Ti,Vector{T}}(hash, new_row_partition, new_col_partition, final_col_indices,
                                new_AT.colptr, new_AT.rowval, new_AT.nzval, local_nrows,
-                               length(final_col_indices), nothing, nothing)
+                               length(final_col_indices), nothing, nothing, new_AT.colptr, new_AT.rowval)
 end
 
 # Convenience: A[row_rng, :] - all columns

src/mumps_factorization.jl

@@ -371,18 +371,23 @@ function _convert_to_mumps_compatible(A::SparseMatrixMPI{T,Ti,AV}) where {T,Ti,A
 
     # Create new SparseMatrixMPI with converted type (CPU Vector)
     # Structural arrays (rowptr, colval, col_indices) stay the same
+    new_rowptr = copy(A.rowptr)
+    new_colval = copy(A.colval)
+    # For CPU, rowptr_target and colval_target are the same as rowptr and colval
     return SparseMatrixMPI{Tinternal,Ti,Vector{Tinternal}}(
         A.structural_hash,
         copy(A.row_partition),
         copy(A.col_partition),
         copy(A.col_indices),
-        copy(A.rowptr),
-        copy(A.colval),
+        new_rowptr,
+        new_colval,
         nzval_converted,
         A.nrows_local,
         A.ncols_compressed,
         nothing,  # Don't copy cached transpose
-        A.cached_symmetric
+        A.cached_symmetric,
+        new_rowptr,  # rowptr_target (same as rowptr for CPU)
+        new_colval   # colval_target (same as colval for CPU)
     )
 end