Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Improve Index Set compat with PartitionedArrays. #4

Open
wants to merge 2 commits into
base: master
Choose a base branch
from
Open

Improve Index Set compat with PartitionedArrays. #4

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
4a029c7
Hotfix IndexSet compat with MPI backend in PartitionedArrays.
termi-official Oct 8, 2022
1e26d47
Type.
termi-official Oct 8, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
51 changes: 33 additions & 18 deletions src/HYPRE.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@
module HYPRE

using MPI: MPI
using PartitionedArrays: IndexRange, MPIData, PSparseMatrix, PVector, PartitionedArrays,
using PartitionedArrays: AbstractIndexSet, MPIData, PSparseMatrix, PVector, PartitionedArrays,
SequentialData, map_parts
using SparseArrays: SparseArrays, SparseMatrixCSC, nnz, nonzeros, nzrange, rowvals
using SparseMatricesCSR: SparseMatrixCSR, colvals, getrowptr
Expand Down Expand Up @@ -305,11 +305,14 @@ end
##################################################

# TODO: This has some duplicated code with to_hypre_data(::SparseMatrixCSC, ilower, iupper)
function Internals.to_hypre_data(A::SparseMatrixCSC, r::IndexRange, c::IndexRange)
@assert r.oid_to_lid isa UnitRange && r.oid_to_lid.start == 1
function Internals.to_hypre_data(A::SparseMatrixCSC, r::AbstractIndexSet, c::AbstractIndexSet)
# @assert r.oid_to_lid isa UnitRange && r.oid_to_lid.start == 1

ilower = r.lid_to_gid[r.oid_to_lid.start]
iupper = r.lid_to_gid[r.oid_to_lid.stop]
lidlower = minimum(r.oid_to_lid)
lidupper = maximum(r.oid_to_lid)

ilower = r.lid_to_gid[lidlower]
iupper = r.lid_to_gid[lidupper]
a_rows = rowvals(A)
a_vals = nonzeros(A)

Expand All @@ -327,7 +330,7 @@ function Internals.to_hypre_data(A::SparseMatrixCSC, r::IndexRange, c::IndexRang
@inbounds for j in 1:size(A, 2)
for i in nzrange(A, j)
row = a_rows[i]
row > r.oid_to_lid.stop && continue # Skip ghost rows
row > lidupper && continue # Skip ghost rows
# grow = r.lid_to_gid[lrow]
ncols[row] += 1
end
Expand All @@ -347,7 +350,7 @@ function Internals.to_hypre_data(A::SparseMatrixCSC, r::IndexRange, c::IndexRang
@inbounds for j in 1:size(A, 2)
for i in nzrange(A, j)
row = a_rows[i]
row > r.oid_to_lid.stop && continue # Skip ghost rows
row > lidupper && continue # Skip ghost rows
k = lastinds[row] += 1
val = a_vals[i]
cols[k] = c.lid_to_gid[j]
Expand All @@ -360,14 +363,17 @@ end
# TODO: Possibly this can be optimized if it is possible to pass overlong vectors to HYPRE.
# At least values should be possible to directly share, but cols needs to translated
# to global ids.
function Internals.to_hypre_data(A::SparseMatrixCSR, r::IndexRange, c::IndexRange)
@assert r.oid_to_lid isa UnitRange && r.oid_to_lid.start == 1
function Internals.to_hypre_data(A::SparseMatrixCSR, r::AbstractIndexSet, c::AbstractIndexSet)
#@assert r.oid_to_lid isa UnitRange && r.oid_to_lid.start == 1

lidlower = minimum(r.oid_to_lid)
lidupper = maximum(r.oid_to_lid)

ilower = r.lid_to_gid[r.oid_to_lid.start]
iupper = r.lid_to_gid[r.oid_to_lid.stop]
ilower = r.lid_to_gid[lidlower]
iupper = r.lid_to_gid[lidupper]
a_cols = colvals(A)
a_vals = nonzeros(A)
nnz = getrowptr(A)[r.oid_to_lid.stop + 1] - 1
nnz = getrowptr(A)[lidupper + 1] - 1

# Initialize the data buffers HYPRE wants
nrows = HYPRE_Int(iupper - ilower + 1) # Total number of rows
Expand Down Expand Up @@ -445,8 +451,11 @@ function HYPREVector(v::PVector)
b = HYPREVector(comm, ilower, iupper)
# Set all the values
map_parts(v.values, v.owned_values, v.rows.partition) do _, vo, vr
ilower_part = vr.lid_to_gid[vr.oid_to_lid.start]
iupper_part = vr.lid_to_gid[vr.oid_to_lid.stop]
lidlower = minimum(vr.oid_to_lid)
lidupper = maximum(vr.oid_to_lid)

ilower_part = vr.lid_to_gid[lidlower]
iupper_part = vr.lid_to_gid[lidupper]

# Option 1: Set all values
nvalues = HYPRE_Int(iupper_part - ilower_part + 1)
Expand Down Expand Up @@ -489,8 +498,11 @@ end
function Base.copy!(dst::PVector{HYPRE_Complex}, src::HYPREVector)
Internals.copy_check(src, dst)
map_parts(dst.values, dst.owned_values, dst.rows.partition) do vv, _, vr
il_src_part = vr.lid_to_gid[vr.oid_to_lid.start]
iu_src_part = vr.lid_to_gid[vr.oid_to_lid.stop]
lidlower = minimum(vr.oid_to_lid)
lidupper = maximum(vr.oid_to_lid)

il_src_part = vr.lid_to_gid[lidlower]
iu_src_part = vr.lid_to_gid[lidupper]
nvalues = HYPRE_Int(iu_src_part - il_src_part + 1)
indices = collect(HYPRE_BigInt, il_src_part:iu_src_part)

Expand All @@ -510,8 +522,11 @@ function Base.copy!(dst::HYPREVector, src::PVector{HYPRE_Complex})
# Re-initialize the vector
@check HYPRE_IJVectorInitialize(dst.ijvector)
map_parts(src.values, src.owned_values, src.rows.partition) do vv, _, vr
ilower_src_part = vr.lid_to_gid[vr.oid_to_lid.start]
iupper_src_part = vr.lid_to_gid[vr.oid_to_lid.stop]
lidlower = minimum(vr.oid_to_lid)
lidupper = maximum(vr.oid_to_lid)

ilower_src_part = vr.lid_to_gid[lidlower]
iupper_src_part = vr.lid_to_gid[lidupper]
nvalues = HYPRE_Int(iupper_src_part - ilower_src_part + 1)
indices = collect(HYPRE_BigInt, ilower_src_part:iupper_src_part)
# TODO: Safe to use vv here? Owned values are always first?
Expand Down