Remove accidental CPU↔GPU conversion helpers for type safety

- Inline GPU staging logic in repartition plans to preserve backend type
- Remove _stage_cpu_to_backend and _stage_cpu_matrix_to_backend helpers
- Inline _mumps_convert_to_backend (MUMPS-only conversion now explicit)
- Fix map_rows to use copy(transpose(...)) preserving GPU arrays
- Fix map_rows type check to use AbstractMatrix instead of Matrix
- Rewrite test_map_rows.jl to use SVector as intended

The type system now catches accidental CPU↔GPU mixing at dispatch time
rather than relying on runtime assertions.

Bump version to 0.1.4

Author: Sébastien Loisel

Project.toml

@@ -1,6 +1,6 @@
 name = "LinearAlgebraMPI"
 uuid = "5bdd2be4-ae34-42ef-8b36-f4c85d48f377"
-version = "0.1.3"
+version = "0.1.4"
 authors = ["S. Loisel"]
 
 [deps]

ext/LinearAlgebraMPIMetalExt.jl

@@ -139,12 +139,13 @@ function LinearAlgebraMPI._array_to_backend(v::Vector{T}, ::Type{<:MtlVector}) w
 end
 
 """
-    _create_output_like(v::LinearAlgebraMPI.VectorMPI{T,<:Vector}, ::Type{<:MtlVector}) where T
+    _convert_vector_to_backend(v::LinearAlgebraMPI.VectorMPI{T,<:Vector}, ::Type{<:MtlVector}) where T
 
-Create a VectorMPI with MtlVector backend from a CPU VectorMPI.
-Used by MUMPS factorization to reconstruct GPU output vectors.
+Convert a CPU VectorMPI to GPU (Metal) backend.
+WARNING: This function exists ONLY for MUMPS, which is a CPU-only solver.
+MUMPS requires GPU→CPU→GPU cycling. Do NOT use this for general operations.
 """
-function LinearAlgebraMPI._create_output_like(v::LinearAlgebraMPI.VectorMPI{T,<:Vector}, ::Type{<:MtlVector}) where T
+function LinearAlgebraMPI._convert_vector_to_backend(v::LinearAlgebraMPI.VectorMPI{T,<:Vector}, ::Type{<:MtlVector}) where T
     return LinearAlgebraMPI.mtl(v)
 end
 

src/LinearAlgebraMPI.jl

@@ -851,7 +851,8 @@ _matrix_to_svectors(M::AbstractMatrix{T}) where {T} = _matrix_to_svectors(Val(si
 # Helper: Convert Vector{SVector} back to Matrix
 function _svectors_to_matrix(v::AbstractVector{SVector{K,T}}) where {K,T}
     # reinterpret as (K, nrows), then transpose to (nrows, K)
-    Matrix(transpose(reinterpret(reshape, T, v)))
+    # Use copy(transpose(...)) to preserve GPU array type (not Matrix(...) which forces CPU)
+    copy(transpose(reinterpret(reshape, T, v)))
 end
 
 # Helper: Convert to SVector representation for map_rows
@@ -941,7 +942,7 @@ function map_rows(f, A...)
     row_partition = first_arg isa VectorMPI ? first_arg.partition : first_arg.row_partition
     hash = compute_partition_hash(row_partition)
 
-    if local_result isa Matrix
+    if local_result isa AbstractMatrix
         return MatrixMPI(
             hash,
             row_partition,

src/dense.jl

@@ -617,7 +617,8 @@ function Base.:*(A::MatrixMPI{T,AM}, x::VectorMPI{T,AV}) where {T,AM,AV}
         # Matrix is on CPU - compute on CPU, then copy to GPU if needed
         y_local_cpu = Vector{T}(undef, local_rows)
         LinearAlgebra.mul!(y_local_cpu, A.A, plan.gathered_cpu)
-        y_v = _create_output_like(x.v, y_local_cpu)
+        # Copy to GPU if input was GPU
+        y_v = x.v isa Vector ? y_local_cpu : copyto!(similar(x.v, local_rows), y_local_cpu)
     else
         # Matrix is on GPU - use GPU gathered buffer directly
         y_v = similar(A.A, local_rows)
@@ -1206,8 +1207,8 @@ function Base.:*(At::Transpose{T,<:MatrixMPI{T}}, x::VectorMPI{T,AV}) where {T,A
     my_col_end = A.col_partition[rank+2] - 1
     local_result_cpu = full_result[my_col_start:my_col_end]
 
-    # Create output with same storage type as input
-    local_result = _create_output_like(x.v, local_result_cpu)
+    # Copy to GPU if input was GPU
+    local_result = x.v isa Vector ? local_result_cpu : copyto!(similar(x.v, length(local_result_cpu)), local_result_cpu)
 
     # Create result vector (partition is immutable, no need to copy)
     y = VectorMPI{T,AV}(
@@ -1470,6 +1471,11 @@ end
 # DenseRepartitionPlan: Repartition a MatrixMPI to a new row partition
 # ============================================================================
 
+# Helper to get CPU copy of matrix data (no-op for CPU, copy for GPU)
+# Used for MPI communication which always requires CPU buffers
+_ensure_cpu(M::Matrix) = M
+_ensure_cpu(M::AbstractMatrix) = Array(M)
+
 """
     DenseRepartitionPlan{T}
 
@@ -1631,28 +1637,31 @@ function DenseRepartitionPlan(A::MatrixMPI{T}, p::Vector{Int}) where T
 end
 
 """
-    execute_plan!(plan::DenseRepartitionPlan{T}, A::MatrixMPI{T}) where T
+    execute_plan!(plan::DenseRepartitionPlan{T}, A::MatrixMPI{T,AM}) where {T,AM}
 
 Execute a dense repartition plan to redistribute rows from A to a new partition.
-Returns a new MatrixMPI with the target row partition.
+Returns a new MatrixMPI with the target row partition, preserving the array type (CPU/GPU).
 """
-function execute_plan!(plan::DenseRepartitionPlan{T}, A::MatrixMPI{T}) where T
+function execute_plan!(plan::DenseRepartitionPlan{T}, A::MatrixMPI{T,AM}) where {T,AM}
     comm = MPI.COMM_WORLD
 
-    # Allocate result
-    result_A = Matrix{T}(undef, plan.result_local_nrows, plan.ncols)
+    # Get CPU version of input for MPI operations
+    A_cpu = _ensure_cpu(A.A)
+
+    # Allocate CPU result buffer (MPI requires CPU)
+    result_cpu = Matrix{T}(undef, plan.result_local_nrows, plan.ncols)
 
-    # Step 1: Local copy
+    # Step 1: Local copy (from CPU staging)
     if !isempty(plan.local_src_range) && !isempty(plan.local_dst_range)
-        result_A[plan.local_dst_range, :] = A.A[plan.local_src_range, :]
+        result_cpu[plan.local_dst_range, :] = A_cpu[plan.local_src_range, :]
     end
 
-    # Step 2: Fill send buffers and send
+    # Step 2: Fill send buffers and send (from CPU staging)
     @inbounds for i in eachindex(plan.send_rank_ids)
         r = plan.send_rank_ids[i]
         row_range = plan.send_row_ranges[i]
         buf = plan.send_bufs[i]
-        buf .= @view A.A[row_range, :]
+        buf .= @view A_cpu[row_range, :]
         plan.send_reqs[i] = MPI.Isend(vec(buf), comm; dest=r, tag=93)
     end
 
@@ -1667,12 +1676,15 @@ function execute_plan!(plan::DenseRepartitionPlan{T}, A::MatrixMPI{T}) where T
     @inbounds for i in eachindex(plan.recv_rank_ids)
         dst_range = plan.recv_dst_ranges[i]
         buf = plan.recv_bufs[i]
-        result_A[dst_range, :] = buf
+        result_cpu[dst_range, :] = buf
     end
 
     MPI.Waitall(plan.send_reqs)
 
-    return MatrixMPI{T}(plan.result_structural_hash, plan.result_row_partition, plan.result_col_partition, result_A)
+    # Copy result back to target array type (GPU if input was GPU)
+    result_A = A.A isa Matrix ? result_cpu : copyto!(similar(A.A, size(result_cpu)...), result_cpu)
+
+    return MatrixMPI{T,typeof(result_A)}(plan.result_structural_hash, plan.result_row_partition, plan.result_col_partition, result_A)
 end
 
 """

src/mumps_factorization.jl

@@ -496,14 +496,12 @@ function _create_output_vector(::Type{Tin}, ::Type{AVin}, n::Int, partition) whe
     return _convert_vector_to_backend(v_dist, AVin)
 end
 
-# Convert VectorMPI to a specific backend type
-function _convert_vector_to_backend(v::VectorMPI{T, AV}, ::Type{AVtarget}) where {T, AV, AVtarget}
-    if AV === AVtarget
-        return v
-    end
-    # Need to convert - handled by mtl/cpu functions in extension
-    return _create_output_like(v, AVtarget)
-end
+# Convert VectorMPI to a specific backend type.
+# WARNING: This function exists ONLY for MUMPS, which is a CPU-only solver.
+# MUMPS requires GPU→CPU→GPU cycling. Do NOT use this for general operations.
+# The base module only defines the identity case (same type).
+# Extensions define CPU→GPU conversions (e.g., Vector → MtlVector).
+_convert_vector_to_backend(v::VectorMPI{T, AV}, ::Type{AV}) where {T, AV} = v
 
 """
     solve(F::MUMPSFactorizationMPI{Tin, AVin, Tinternal}, b::VectorMPI) where {Tin, AVin, Tinternal}
@@ -600,9 +598,6 @@ end
 # Fallback for CPU arrays
 _array_to_backend(v::Vector{T}, ::Type{<:Vector}) where T = v
 
-# Fallback for CPU VectorMPI (no conversion needed)
-_create_output_like(v::VectorMPI{T,<:Vector}, ::Type{<:Vector}) where T = v
-
 """
     Base.:\\(F::MUMPSFactorizationMPI, b::VectorMPI)
 

src/sparse.jl

@@ -1224,7 +1224,7 @@ function Base.:*(A::SparseMatrixMPI{T,Ti,AV}, B::SparseMatrixMPI{T,Ti,AV}) where
     nzval = similar(A.nzval, nnz_result)
 
     # Adapt gathered B values (plan.AT.nzval is CPU) to A's backend
-    AT_nzval = _create_output_like(A.nzval, plan.AT.nzval)
+    AT_nzval = A.nzval isa Vector ? plan.AT.nzval : copyto!(similar(A.nzval, length(plan.AT.nzval)), plan.AT.nzval)
 
     # Execute symbolic multiplication (unified kernel)
     _execute_symbolic_multiply!(nzval, plan, AT_nzval, A.nzval)
@@ -1589,11 +1589,9 @@ function Base.:+(A::SparseMatrixMPI{T,Ti,AV}, B::SparseMatrixMPI{T,Ti,AV}) where
     nnz_result = plan.colptr[end] - 1
     nzval = similar(A.nzval, nnz_result)
 
-    # Adapt B_repart.nzval to A's backend (repartition always returns CPU)
-    B_nzval = _create_output_like(A.nzval, B_repart.nzval)
-
     # Execute addition directly into result buffer (unified kernel)
-    execute_addition!(nzval, plan, A.nzval, B_nzval)
+    # B_repart.nzval has same backend as A.nzval (repartition preserves backend)
+    execute_addition!(nzval, plan, A.nzval, B_repart.nzval)
 
     # Extract CSR components directly (rowptr=colptr, colval=rowval from plan)
     nrows_local = length(plan.colptr) - 1
@@ -1631,11 +1629,9 @@ function Base.:-(A::SparseMatrixMPI{T,Ti,AV}, B::SparseMatrixMPI{T,Ti,AV}) where
     nnz_result = plan.colptr[end] - 1
     nzval = similar(A.nzval, nnz_result)
 
-    # Adapt B_repart.nzval to A's backend (repartition always returns CPU)
-    B_nzval = _create_output_like(A.nzval, B_repart.nzval)
-
     # Execute subtraction directly into result buffer (unified kernel)
-    execute_subtraction!(nzval, plan, A.nzval, B_nzval)
+    # B_repart.nzval has same backend as A.nzval (repartition preserves backend)
+    execute_subtraction!(nzval, plan, A.nzval, B_repart.nzval)
 
     # Extract CSR components directly (rowptr=colptr, colval=rowval from plan)
     nrows_local = length(plan.colptr) - 1
@@ -1975,7 +1971,7 @@ function execute_plan!(plan::TransposePlan{T,Ti}, A::SparseMatrixMPI{T,Ti,AV}) w
     ncols_compressed = length(plan.col_indices)
 
     # Copy result to GPU if input was GPU
-    result_nzval = _create_output_like(A.nzval, compressed_result_AT.nzval)
+    result_nzval = A.nzval isa Vector ? compressed_result_AT.nzval : copyto!(similar(A.nzval, length(compressed_result_AT.nzval)), compressed_result_AT.nzval)
 
     # Convert structure arrays to target backend
     rowptr_target = _to_target_backend(compressed_result_AT.colptr, AV)
@@ -2260,8 +2256,8 @@ function Base.:*(A::SparseMatrixMPI{T,Ti,AV}, x::VectorMPI{T,AVX}) where {T,Ti,A
     # Get gathered data on CPU (MPI communication is always CPU)
     gathered_cpu = _gathered_cpu_buffer(plan)
 
-    # Adapt gathered data to A's backend and allocate result
-    gathered = _create_output_like(A.nzval, gathered_cpu)
+    # Adapt gathered data to A's backend
+    gathered = A.nzval isa Vector ? gathered_cpu : copyto!(similar(A.nzval, length(gathered_cpu)), gathered_cpu)
     y_local = similar(A.nzval, local_rows)
 
     # Unified SpMV kernel using pre-computed target arrays (works on CPU or GPU)

src/vectors.jl

@@ -111,12 +111,12 @@ function Adapt.adapt_structure(to, v::VectorMPI{T,AV}) where {T,AV}
 end
 
 # Helper to create output vector with same storage type as reference
-# Used by sparse and dense A*x operations
+# STRICT: Only allows same-type matching. Mismatched types will error.
+# This prevents accidental CPU↔GPU mixing.
 _create_output_like(::Vector{T}, result::Vector{T}) where T = result
-_create_output_like(ref::AV, result::Vector{T}) where {T,AV<:AbstractVector{T}} =
-    copyto!(similar(ref, length(result)), result)
-# When result is already the same type as reference, return as-is (e.g., both GPU)
 _create_output_like(::AV, result::AV) where {T,AV<:AbstractVector{T}} = result
+# Note: No catch-all method for (GPU, CPU) - this is intentional to catch bugs.
+# MPI staging (CPU→GPU) is inlined where needed with explicit ternary checks.
 
 # Helper to get CPU copy of vector data (no-op for CPU, copy for GPU)
 # Used for MPI communication which always requires CPU buffers
@@ -503,31 +503,34 @@ function VectorRepartitionPlan(x::VectorMPI{T}, p::Vector{Int}) where T
 end
 
 """
-    execute_plan!(plan::VectorRepartitionPlan{T}, x::VectorMPI{T}) where T
+    execute_plan!(plan::VectorRepartitionPlan{T}, x::VectorMPI{T,AV}) where {T,AV}
 
 Execute a vector repartition plan to redistribute elements from x to a new partition.
-Returns a new VectorMPI with the target partition.
+Returns a new VectorMPI with the target partition, preserving the array type (CPU/GPU).
 """
-function execute_plan!(plan::VectorRepartitionPlan{T}, x::VectorMPI{T}) where T
+function execute_plan!(plan::VectorRepartitionPlan{T}, x::VectorMPI{T,AV}) where {T,AV}
     comm = MPI.COMM_WORLD
 
-    # Allocate result
-    result_v = Vector{T}(undef, plan.result_local_size)
+    # Get CPU version of input for MPI operations
+    x_cpu = _ensure_cpu(x.v)
+
+    # Allocate CPU result buffer (MPI requires CPU)
+    result_cpu = Vector{T}(undef, plan.result_local_size)
 
-    # Step 1: Local copy
+    # Step 1: Local copy (from CPU staging)
     if !isempty(plan.local_src_range)
         @inbounds for (i, src_i) in enumerate(plan.local_src_range)
-            result_v[plan.local_dst_offset + i - 1] = x.v[src_i]
+            result_cpu[plan.local_dst_offset + i - 1] = x_cpu[src_i]
         end
     end
 
-    # Step 2: Fill send buffers and send
+    # Step 2: Fill send buffers and send (from CPU staging)
     @inbounds for i in eachindex(plan.send_rank_ids)
         r = plan.send_rank_ids[i]
         range = plan.send_ranges[i]
         buf = plan.send_bufs[i]
         for (k, src_k) in enumerate(range)
-            buf[k] = x.v[src_k]
+            buf[k] = x_cpu[src_k]
         end
         plan.send_reqs[i] = MPI.Isend(buf, comm; dest=r, tag=92)
     end
@@ -544,13 +547,16 @@ function execute_plan!(plan::VectorRepartitionPlan{T}, x::VectorMPI{T}) where T
         offset = plan.recv_offsets[i]
         buf = plan.recv_bufs[i]
         for k in eachindex(buf)
-            result_v[offset + k - 1] = buf[k]
+            result_cpu[offset + k - 1] = buf[k]
         end
     end
 
     MPI.Waitall(plan.send_reqs)
 
-    return VectorMPI{T}(plan.result_partition_hash, plan.result_partition, result_v)
+    # Copy result back to target array type (GPU if input was GPU)
+    result_v = x.v isa Vector ? result_cpu : copyto!(similar(x.v, length(result_cpu)), result_cpu)
+
+    return VectorMPI{T,typeof(result_v)}(plan.result_partition_hash, plan.result_partition, result_v)
 end
 
 """