JuliaSmoothOptimizers / KrylovPreconditioners.jl

The ultimate collection of preconditioners
Other
7 stars 1 forks source link

[GPU] Add kernels to scale CSC and CSR matrices #34

Open amontoison opened 8 months ago

amontoison commented 8 months ago
function scaling_csc!(A::SparseMatrixCSC{T}, s::Vector{T}) where T
  n = length(s)
  @assert size(A,2) == n
  for j = 1:n
    acc = zero(T)
    @inbounds @simd for i = A.colptr[j]:(A.colptr[j + 1] - 1)
      acc += A.nzval[i] * A.nzval[i]
    end
    s[j] = sqrt(acc)
    if acc ≠ zero(T)
      @inbounds @simd for i = A.colptr[j]:(A.colptr[j + 1] - 1)
        A.nzval[i] /= s[j]
      end
    end
  end
end
function scaling_csc!(A::CuSparseMatrixCSC{T}, s::CuVector{T}) where T
  n = length(s)
  @assert size(A,2) == n
  for j = 1:n
    acc = zero(T)
    @inbounds @simd for i = A.colPtr[j]:(A.colPtr[j + 1] - 1)
      acc += A.nzVal[i] * A.nzVal[i]
    end
    s[j] = sqrt(acc)
    if acc ≠ zero(T)
      @inbounds @simd for i = A.colPtr[j]:(A.colPtr[j + 1] - 1)
        A.nzVal[i] /= s[j]
      end
    end
  end
end
function scaling_csr!(A::CuSparseMatrixCSR{T}, s::CuVector{T}) where T
  m = length(s)
  @assert size(A,1) == m
  for j = 1:m
    acc = zero(T)
    @inbounds @simd for i = A.rowPtr[j]:(A.rowPtr[j + 1] - 1)
      acc += A.nzVal[i] * A.nzVal[i]
    end
    s[j] = sqrt(acc)
    if acc ≠ zero(T)
      @inbounds @simd for i = A.rowPtr[j]:(A.rowPtr[j + 1] - 1)
        A.nzVal[i] /= s[j]
      end
    end
  end
end

cc @michel2323