From 497bfa913a2f635595c9842be10f52a23b3296c2 Mon Sep 17 00:00:00 2001 From: Fredrik Ekre Date: Sun, 29 Sep 2024 01:12:57 +0200 Subject: [PATCH] Move PartitionedArrays support to an extension --- .github/workflows/Check.yml | 31 +++- CHANGELOG.md | 2 + Project.toml | 9 +- docs/Manifest.toml | 102 +---------- ext/HYPREPartitionedArrays.jl | 324 ++++++++++++++++++++++++++++++++++ src/HYPRE.jl | 301 +------------------------------ src/solvers.jl | 18 -- 7 files changed, 371 insertions(+), 416 deletions(-) create mode 100644 ext/HYPREPartitionedArrays.jl diff --git a/.github/workflows/Check.yml b/.github/workflows/Check.yml index 03db225..df6ebcf 100644 --- a/.github/workflows/Check.yml +++ b/.github/workflows/Check.yml @@ -22,18 +22,37 @@ jobs: # with: # version: '1' - uses: julia-actions/cache@v2 - - uses: julia-actions/julia-buildpkg@v1 + # - uses: julia-actions/julia-buildpkg@v1 - name: Install dependencies - shell: julia {0} + shell: julia --project=@explicit-imports {0} run: | - # Add ExplicitImports.jl + # Add ExplicitImports.jl and packages that HYPRE has extensions for using Pkg - Pkg.add([PackageSpec(name = "ExplicitImports", version = "1.9")]) + Pkg.develop([ + PackageSpec(name = "HYPRE", path = pwd()), + ]) + Pkg.add([ + PackageSpec(name = "ExplicitImports", version = "1.9"), + PackageSpec(name = "PartitionedArrays", version = "0.5"), + ]) - name: ExplicitImports.jl code checks - shell: julia --project {0} + shell: julia --project=@explicit-imports {0} run: | - using HYPRE, ExplicitImports + using HYPRE, ExplicitImports, PartitionedArrays + # Check HYPRE check_no_implicit_imports(HYPRE) check_no_stale_explicit_imports(HYPRE) check_all_qualified_accesses_via_owners(HYPRE) check_no_self_qualified_accesses(HYPRE) + # Check extension modules + for ext in (:HYPREPartitionedArrays,) + extmod = Base.get_extension(HYPRE, ext) + if extmod !== nothing + check_no_implicit_imports(extmod) + check_no_stale_explicit_imports(extmod) + check_all_qualified_accesses_via_owners(extmod) + check_no_self_qualified_accesses(extmod) + else + @warn "$(ext) extension not available." + end + end diff --git a/CHANGELOG.md b/CHANGELOG.md index 35a1aa2..43b4e39 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,6 +9,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Changed - PartitionedArrays.jl dependency upgraded from release series 0.3.x to release series 0.5.x. ([#17], [#18]) + - PartitionedArrays.jl support is now moved to a package extension. ([#23]) - CEnum.jl dependency upgraded to release series 0.5.x (release series 0.4.x still allowed). ([#17], [#18]) @@ -90,3 +91,4 @@ Initial release of HYPRE.jl. [#16]: https://github.com/fredrikekre/HYPRE.jl/issues/16 [#17]: https://github.com/fredrikekre/HYPRE.jl/issues/17 [#18]: https://github.com/fredrikekre/HYPRE.jl/issues/18 +[#23]: https://github.com/fredrikekre/HYPRE.jl/issues/23 diff --git a/Project.toml b/Project.toml index 9563442..b7be1e5 100644 --- a/Project.toml +++ b/Project.toml @@ -11,6 +11,12 @@ PartitionedArrays = "5a9dfac6-5c52-46f7-8278-5e2210713be9" SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf" SparseMatricesCSR = "a0a7dd2c-ebf4-11e9-1f05-cf50bc540ca1" +[weakdeps] +PartitionedArrays = "5a9dfac6-5c52-46f7-8278-5e2210713be9" + +[extensions] +HYPREPartitionedArrays = "PartitionedArrays" + [compat] CEnum = "0.4, 0.5" MPI = "0.19, 0.20" @@ -20,7 +26,8 @@ julia = "1.6" [extras] LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" +PartitionedArrays = "5a9dfac6-5c52-46f7-8278-5e2210713be9" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" [targets] -test = ["LinearAlgebra", "Test"] +test = ["LinearAlgebra", "PartitionedArrays", "Test"] diff --git a/docs/Manifest.toml b/docs/Manifest.toml index cea96be..3e402c2 100644 --- a/docs/Manifest.toml +++ b/docs/Manifest.toml @@ -14,38 +14,16 @@ git-tree-sha1 = "faa260e4cb5aba097a73fab382dd4b5819d8ec8c" uuid = "1520ce14-60c1-5f80-bbc7-55ef81b5835c" version = "0.4.4" -[[deps.Adapt]] -deps = ["LinearAlgebra", "Requires"] -git-tree-sha1 = "76289dc51920fdc6e0013c872ba9551d54961c24" -uuid = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" -version = "3.6.2" -weakdeps = ["StaticArrays"] - - [deps.Adapt.extensions] - AdaptStaticArraysExt = "StaticArrays" - [[deps.ArgTools]] uuid = "0dad84c5-d112-42e6-8d28-ef12dabb789f" version = "1.1.1" -[[deps.ArrayLayouts]] -deps = ["FillArrays", "LinearAlgebra", "SparseArrays"] -git-tree-sha1 = "1d9e98721e71dcf4db5a7d34f55d8aa07c43468f" -uuid = "4c555306-a7a7-4459-81d9-ec55ddd5c99a" -version = "1.0.6" - [[deps.Artifacts]] uuid = "56f22d72-fd6d-98f1-02f0-08ddc0907c33" [[deps.Base64]] uuid = "2a0f44e3-6c83-55bd-87e4-b1978d98bd5f" -[[deps.BlockArrays]] -deps = ["ArrayLayouts", "FillArrays", "LinearAlgebra"] -git-tree-sha1 = "bed8cfec0c348753a79d915cc82999c395299297" -uuid = "8e7c35d0-a365-5155-bbbb-fb81a777f24e" -version = "0.16.33" - [[deps.CEnum]] git-tree-sha1 = "eb4cb44a499229b3b8426dcfb5dd85333951ff90" uuid = "fa961155-64e5-5f13-b03f-caf6b980ea82" @@ -56,12 +34,6 @@ git-tree-sha1 = "e579c6157598169ad4ef17263bdf3452b4a3e316" uuid = "5217a498-cd5d-4ec6-b8c2-9b85a09b6e3e" version = "1.1.0" -[[deps.CircularArrays]] -deps = ["OffsetArrays"] -git-tree-sha1 = "61bc114e595167090b4cbcb7305ddeacd4274f16" -uuid = "7a955b69-7140-5f4e-a0ed-f168c5e2e749" -version = "1.3.2" - [[deps.CodecZlib]] deps = ["TranscodingStreams", "Zlib_jll"] git-tree-sha1 = "bce6804e5e6044c6daab27bb533d1295e4a2e759" @@ -77,12 +49,6 @@ version = "1.1.1+0" deps = ["Printf"] uuid = "ade2ca70-3891-5945-98fb-dc099432e06a" -[[deps.Distances]] -deps = ["LinearAlgebra", "SparseArrays", "Statistics", "StatsAPI"] -git-tree-sha1 = "49eba9ad9f7ead780bfb7ee319f962c811c6d3b2" -uuid = "b4f34e82-e78d-54a5-968a-f98e89d6e8f7" -version = "0.10.8" - [[deps.Distributed]] deps = ["Random", "Serialization", "Sockets"] uuid = "8ba89e20-285c-5b6f-9357-94700520ee1b" @@ -113,12 +79,6 @@ version = "2.6.2+0" [[deps.FileWatching]] uuid = "7b1f6079-737a-58dc-b8bc-7a2ca5c1b5ee" -[[deps.FillArrays]] -deps = ["LinearAlgebra", "Random", "SparseArrays", "Statistics"] -git-tree-sha1 = "ed569cb9e7e3590d5ba884da7edc50216aac5811" -uuid = "1a297f60-69ca-5386-bcde-b61e274b549b" -version = "1.1.0" - [[deps.Git]] deps = ["Git_jll"] git-tree-sha1 = "04eff47b1354d702c3a85e8ab23d539bb7d5957e" @@ -132,11 +92,17 @@ uuid = "f8c6e375-362e-5223-8a59-34ff63f689eb" version = "2.46.2+0" [[deps.HYPRE]] -deps = ["CEnum", "HYPRE_jll", "Libdl", "MPI", "PartitionedArrays", "SparseArrays", "SparseMatricesCSR"] +deps = ["CEnum", "HYPRE_jll", "Libdl", "MPI", "SparseArrays", "SparseMatricesCSR"] path = ".." uuid = "b5ffcf37-a2bd-41ab-a3da-4bd9bc8ad771" version = "1.5.0" + [deps.HYPRE.extensions] + HYPREPartitionedArrays = "PartitionedArrays" + + [deps.HYPRE.weakdeps] + PartitionedArrays = "5a9dfac6-5c52-46f7-8278-5e2210713be9" + [[deps.HYPRE_jll]] deps = ["Artifacts", "JLLWrappers", "LAPACK_jll", "LazyArtifacts", "Libdl", "MPICH_jll", "MPIPreferences", "MPItrampoline_jll", "MicrosoftMPI_jll", "OpenBLAS_jll", "OpenMPI_jll", "Pkg", "TOML"] git-tree-sha1 = "b77d3eca75f8442e034ccf415c87405a49e77985" @@ -153,12 +119,6 @@ version = "0.2.3" deps = ["Markdown"] uuid = "b77e0a4c-d291-57a0-90e8-8db25a27a240" -[[deps.IterativeSolvers]] -deps = ["LinearAlgebra", "Printf", "Random", "RecipesBase", "SparseArrays"] -git-tree-sha1 = "1169632f425f79429f245113b775a0e3d121457c" -uuid = "42fd0dbc-a981-5370-80f2-aaf504508153" -version = "0.9.2" - [[deps.JLLWrappers]] deps = ["Preferences"] git-tree-sha1 = "abc9885a7ca2052a736a600f7fa66209f96506e1" @@ -290,12 +250,6 @@ version = "2023.1.10" uuid = "ca575930-c2e3-43a9-ace4-1e988b2c1908" version = "1.2.0" -[[deps.OffsetArrays]] -deps = ["Adapt"] -git-tree-sha1 = "82d7c9e310fe55aa54996e6f7f94674e2a38fcb4" -uuid = "6fe1bfb0-de20-5000-8ca7-80f57d26f881" -version = "1.12.9" - [[deps.OpenBLAS_jll]] deps = ["Artifacts", "CompilerSupportLibraries_jll", "Libdl"] uuid = "4536629a-c528-5b80-bd46-f80d51c5b363" @@ -324,12 +278,6 @@ git-tree-sha1 = "a5aef8d4a6e8d81f171b2bd4be5265b01384c74c" uuid = "69de0a69-1ddd-5017-9359-2bf0b02dc9f0" version = "2.5.10" -[[deps.PartitionedArrays]] -deps = ["BlockArrays", "CircularArrays", "Distances", "FillArrays", "IterativeSolvers", "LinearAlgebra", "MPI", "Printf", "Random", "SparseArrays", "SparseMatricesCSR", "StaticArrays"] -git-tree-sha1 = "912e48995e001f7b24d425fb1eebbb97e8c3cb09" -uuid = "5a9dfac6-5c52-46f7-8278-5e2210713be9" -version = "0.5.4" - [[deps.Pkg]] deps = ["Artifacts", "Dates", "Downloads", "FileWatching", "LibGit2", "Libdl", "Logging", "Markdown", "Printf", "REPL", "Random", "SHA", "Serialization", "TOML", "Tar", "UUIDs", "p7zip_jll"] uuid = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f" @@ -359,12 +307,6 @@ uuid = "3fa0cd96-eef1-5676-8a61-b3b8758bbffb" deps = ["SHA"] uuid = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" -[[deps.RecipesBase]] -deps = ["PrecompileTools"] -git-tree-sha1 = "5c3d09cc4f31f5fc6af001c250bf1278733100ff" -uuid = "3cdcf5f2-1ef4-517c-9805-6587b60abb01" -version = "1.3.4" - [[deps.RegistryInstances]] deps = ["LazilyInitializedFields", "Pkg", "TOML", "Tar"] git-tree-sha1 = "ffd19052caf598b8653b99404058fce14828be51" @@ -398,36 +340,6 @@ git-tree-sha1 = "38677ca58e80b5cad2382e5a1848f93b054ad28d" uuid = "a0a7dd2c-ebf4-11e9-1f05-cf50bc540ca1" version = "0.6.7" -[[deps.StaticArrays]] -deps = ["LinearAlgebra", "PrecompileTools", "Random", "StaticArraysCore"] -git-tree-sha1 = "eeafab08ae20c62c44c8399ccb9354a04b80db50" -uuid = "90137ffa-7385-5640-81b9-e52037218182" -version = "1.9.7" - - [deps.StaticArrays.extensions] - StaticArraysChainRulesCoreExt = "ChainRulesCore" - StaticArraysStatisticsExt = "Statistics" - - [deps.StaticArrays.weakdeps] - ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" - Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2" - -[[deps.StaticArraysCore]] -git-tree-sha1 = "192954ef1208c7019899fbf8049e717f92959682" -uuid = "1e83bf80-4336-4d27-bf5d-d5a4f845583c" -version = "1.4.3" - -[[deps.Statistics]] -deps = ["LinearAlgebra", "SparseArrays"] -uuid = "10745b16-79ce-11e8-11f9-7d13ad32a3b2" -version = "1.10.0" - -[[deps.StatsAPI]] -deps = ["LinearAlgebra"] -git-tree-sha1 = "45a7769a04a3cf80da1c1c7c60caf932e6f4c9f7" -uuid = "82ae8749-77ed-4fe6-ae5f-f523153014b0" -version = "1.6.0" - [[deps.SuiteSparse]] deps = ["Libdl", "LinearAlgebra", "Serialization", "SparseArrays"] uuid = "4607b0f0-06f3-5cda-b6b1-a6196a1729e9" diff --git a/ext/HYPREPartitionedArrays.jl b/ext/HYPREPartitionedArrays.jl new file mode 100644 index 0000000..4fed004 --- /dev/null +++ b/ext/HYPREPartitionedArrays.jl @@ -0,0 +1,324 @@ +module HYPREPartitionedArrays + +using HYPRE.LibHYPRE: @check, HYPRE_BigInt, HYPRE_Complex, HYPRE_IJMatrixSetValues, + HYPRE_IJVectorGetValues, HYPRE_IJVectorInitialize, HYPRE_IJVectorSetValues, HYPRE_Int +using HYPRE: HYPRE, HYPREMatrix, HYPRESolver, HYPREVector, Internals +using MPI: MPI +using PartitionedArrays: PartitionedArrays, AbstractLocalIndices, MPIArray, PSparseMatrix, + PVector, SplitMatrix, ghost_to_global, local_values, own_to_global, own_values, + partition +using SparseArrays: SparseArrays, SparseMatrixCSC, nonzeros, nzrange, rowvals +using SparseMatricesCSR: SparseMatrixCSR, colvals + +################################################## +# PartitionedArrays.PSparseMatrix -> HYPREMatrix # +################################################## + +function subarray_unsafe_supported() + # Wrapping of SubArrays as raw pointers may or may not be supported + # depending on the Julia version. If this is not supported, we have to fall + # back to allocation of an intermediate buffer. This logic can be removed if + # HYPRE.jl drops support for Julia < 1.9. + return VERSION >= v"1.9.0" +end + +function Internals.to_hypre_data( + A::SplitMatrix{<:SparseMatrixCSC}, r::AbstractLocalIndices, c::AbstractLocalIndices + ) + # Own/ghost to global index mappings + own_to_global_row = own_to_global(r) + own_to_global_col = own_to_global(c) + ghost_to_global_col = ghost_to_global(c) + + # HYPRE requires contiguous row indices + ilower = own_to_global_row[1] + iupper = own_to_global_row[end] + @assert iupper - ilower + 1 == length(own_to_global_row) + + # Extract sparse matrices from the SplitMatrix. We are only interested in the owned + # rows, so only consider own-own and own-ghost blocks. + Aoo = A.blocks.own_own::SparseMatrixCSC + Aoo_rows = rowvals(Aoo) + Aoo_vals = nonzeros(Aoo) + Aog = A.blocks.own_ghost::SparseMatrixCSC + Aog_rows = rowvals(Aog) + Aog_vals = nonzeros(Aog) + @assert size(Aoo, 1) == size(Aog, 1) == length(own_to_global_row) + + # Initialize the data buffers HYPRE wants + nrows = HYPRE_Int(length(own_to_global_row)) # Total number of rows + ncols = zeros(HYPRE_Int, nrows) # Number of colums for each row + rows = collect(HYPRE_BigInt, ilower:iupper) # The row indices + # cols = Vector{HYPRE_BigInt}(undef, nnz) # The column indices + # values = Vector{HYPRE_Complex}(undef, nnz) # The values + + # First pass to count nnz per row (note that global column indices and column + # permutation doesn't matter for this pass) + @inbounds for own_col in 1:size(Aoo, 2) + for k in nzrange(Aoo, own_col) + own_row = Aoo_rows[k] + ncols[own_row] += 1 + end + end + @inbounds for ghost_col in 1:size(Aog, 2) + for k in nzrange(Aog, ghost_col) + own_row = Aog_rows[k] + ncols[own_row] += 1 + end + end + + # Initialize remaining buffers now that nnz is known + nnz = sum(ncols) + cols = Vector{HYPRE_BigInt}(undef, nnz) + values = Vector{HYPRE_Complex}(undef, nnz) + + # Keep track of the last index used for every row + lastinds = zeros(Int, nrows) + cumsum!((@view lastinds[2:end]), (@view ncols[1:end-1])) + + # Second pass to populate the output. Here we need to map column + # indices from own/ghost to global + @inbounds for own_col in 1:size(Aoo, 2) + for k in nzrange(Aoo, own_col) + own_row = Aoo_rows[k] + i = lastinds[own_row] += 1 + values[i] = Aoo_vals[k] + cols[i] = own_to_global_col[own_col] + end + end + @inbounds for ghost_col in 1:size(Aog, 2) + for k in nzrange(Aog, ghost_col) + own_row = Aog_rows[k] + i = lastinds[own_row] += 1 + values[i] = Aog_vals[k] + cols[i] = ghost_to_global_col[ghost_col] + end + end + + # Sanity checks and return + @assert nrows == length(ncols) == length(rows) + return nrows, ncols, rows, cols, values +end + +function Internals.to_hypre_data( + A::SplitMatrix{<:SparseMatrixCSR}, r::AbstractLocalIndices, c::AbstractLocalIndices + ) + # Own/ghost to global index mappings + own_to_global_row = own_to_global(r) + own_to_global_col = own_to_global(c) + ghost_to_global_col = ghost_to_global(c) + + # HYPRE requires contiguous row indices + ilower = own_to_global_row[1] + iupper = own_to_global_row[end] + @assert iupper - ilower + 1 == length(own_to_global_row) + + # Extract sparse matrices from the SplitMatrix. We are only interested in the owned + # rows, so only consider own-own and own-ghost blocks. + Aoo = A.blocks.own_own::SparseMatrixCSR + Aoo_cols = colvals(Aoo) + Aoo_vals = nonzeros(Aoo) + Aog = A.blocks.own_ghost::SparseMatrixCSR + Aog_cols = colvals(Aog) + Aog_vals = nonzeros(Aog) + @assert size(Aoo, 1) == size(Aog, 1) == length(own_to_global_row) + + # Initialize the data buffers HYPRE wants + nnz = SparseArrays.nnz(Aoo) + SparseArrays.nnz(Aog) + nrows = HYPRE_Int(iupper - ilower + 1) # Total number of rows + ncols = zeros(HYPRE_Int, nrows) # Number of columns for each row + rows = collect(HYPRE_BigInt, ilower:iupper) # The row indices + cols = Vector{HYPRE_BigInt}(undef, nnz) # The column indices + values = Vector{HYPRE_Complex}(undef, nnz) # The values + + # For CSR we only need a single pass to over the owned rows to collect everything + i = 0 + for own_row in 1:size(Aoo, 1) + nzro = nzrange(Aoo, own_row) + nzrg = nzrange(Aog, own_row) + ncols[own_row] = length(nzro) + length(nzrg) + for k in nzro + i += 1 + own_col = Aoo_cols[k] + cols[i] = own_to_global_col[own_col] + values[i] = Aoo_vals[k] + end + for k in nzrg + i += 1 + ghost_col = Aog_cols[k] + cols[i] = ghost_to_global_col[ghost_col] + values[i] = Aog_vals[k] + end + end + + # Sanity checks and return + @assert nnz == i + @assert nrows == length(ncols) == length(rows) + return nrows, ncols, rows, cols, values +end + +function Internals.get_comm(A::Union{PSparseMatrix{<:Any,<:M}, PVector{<:Any,<:M}}) where M <: MPIArray + return partition(A).comm +end + +Internals.get_comm(_::Union{PSparseMatrix,PVector}) = MPI.COMM_SELF + +function Internals.get_proc_rows(A::Union{PSparseMatrix, PVector}) + if A isa PVector + r = A.index_partition + else + r = A.row_partition + end + ilower::HYPRE_BigInt = typemax(HYPRE_BigInt) + iupper::HYPRE_BigInt = typemin(HYPRE_BigInt) + map(r) do a + # This is a map over the local process' owned indices. For MPI it will + # be a single value but for DebugArray / Array it will have multiple + # values. + o_to_g = own_to_global(a) + ilower_part = o_to_g[1] + iupper_part = o_to_g[end] + ilower = min(ilower, convert(HYPRE_BigInt, ilower_part)) + iupper = max(iupper, convert(HYPRE_BigInt, iupper_part)) + end + return ilower, iupper +end + +function HYPREMatrix(B::PSparseMatrix) + # Use the same communicator as the matrix + comm = Internals.get_comm(B) + # Fetch rows owned by this process + ilower, iupper = Internals.get_proc_rows(B) + # Create the IJ matrix + A = HYPREMatrix(comm, ilower, iupper) + # Set all the values + map(local_values(B), B.row_partition, B.col_partition) do Bv, Br, Bc + nrows, ncols, rows, cols, values = Internals.to_hypre_data(Bv, Br, Bc) + @check HYPRE_IJMatrixSetValues(A, nrows, ncols, rows, cols, values) + return nothing + end + # Finalize + Internals.assemble_matrix(A) + return A +end + +############################################ +# PartitionedArrays.PVector -> HYPREVector # +############################################ + +function HYPREVector(v::PVector) + # Use the same communicator as the matrix + comm = Internals.get_comm(v) + # Fetch rows owned by this process + ilower, iupper = Internals.get_proc_rows(v) + # Create the IJ vector + b = HYPREVector(comm, ilower, iupper) + # Set all the values + map(own_values(v), v.index_partition) do vo, vr + o_to_g = own_to_global(vr) + + ilower_part = o_to_g[1] + iupper_part = o_to_g[end] + + # Option 1: Set all values + nvalues = HYPRE_Int(iupper_part - ilower_part + 1) + indices = collect(HYPRE_BigInt, ilower_part:iupper_part) + # TODO: Could probably just pass the full vector even if it is too long + # values = convert(Vector{HYPRE_Complex}, vv) + values = collect(HYPRE_Complex, vo) + + # # Option 2: Set only non-zeros + # indices = HYPRE_BigInt[] + # values = HYPRE_Complex[] + # for (i, vi) in zip(ilower_part:iupper_part, vo) + # if !iszero(vi) + # push!(indices, i) + # push!(values, vi) + # end + # end + # nvalues = length(indices) + + @check HYPRE_IJVectorSetValues(b, nvalues, indices, values) + return nothing + end + # Finalize + Internals.assemble_vector(b) + return b +end + +function copy_check(dst::HYPREVector, src::PVector) + il_dst, iu_dst = Internals.get_proc_rows(dst) + il_src, iu_src = Internals.get_proc_rows(src) + if il_dst != il_src && iu_dst != iu_src + # TODO: Why require this? + throw(ArgumentError( + "row owner mismatch between dst ($(il_dst:iu_dst)) and src ($(il_src:iu_src))" + )) + end +end + +# TODO: Other eltypes could be support by using a intermediate buffer +function Base.copy!(dst::PVector{<:AbstractVector{HYPRE_Complex}}, src::HYPREVector) + copy_check(src, dst) + map(own_values(dst), dst.index_partition) do ov, vr + o_to_g = own_to_global(vr) + il_src_part = o_to_g[1] + iu_src_part = o_to_g[end] + nvalues = HYPRE_Int(iu_src_part - il_src_part + 1) + indices = collect(HYPRE_BigInt, il_src_part:iu_src_part) + if subarray_unsafe_supported() + values = ov + else + values = collect(HYPRE_Complex, ov) + end + @check HYPRE_IJVectorGetValues(src, nvalues, indices, values) + if !subarray_unsafe_supported() + @. ov = values + end + end + return dst +end + +function Base.copy!(dst::HYPREVector, src::PVector{<:AbstractVector{HYPRE_Complex}}) + copy_check(dst, src) + # Re-initialize the vector + @check HYPRE_IJVectorInitialize(dst) + map(own_values(src), src.index_partition) do ov, vr + o_to_g = own_to_global(vr) + ilower_src_part = o_to_g[1] + iupper_src_part = o_to_g[end] + nvalues = HYPRE_Int(iupper_src_part - ilower_src_part + 1) + indices = collect(HYPRE_BigInt, ilower_src_part:iupper_src_part) + if subarray_unsafe_supported() + values = ov + else + values = collect(HYPRE_Complex, ov) + end + @check HYPRE_IJVectorSetValues(dst, nvalues, indices, values) + end + # TODO: It shouldn't be necessary to assemble here since we only set owned rows (?) + # @check HYPRE_IJVectorAssemble(dst) + # TODO: Necessary to recreate the ParVector? Running some examples it seems like it is + # not needed. + return dst +end + +###################################### +# PartitionedArrays solver interface # +###################################### + +# TODO: Would it be useful with a method that copied the solution to b instead? + +function HYPRE.solve(solver::HYPRESolver, A::PSparseMatrix, b::PVector) + hypre_x = HYPRE.solve(solver, HYPREMatrix(A), HYPREVector(b)) + x = copy!(similar(b, HYPRE_Complex), hypre_x) + return x +end +function HYPRE.solve!(solver::HYPRESolver, x::PVector, A::PSparseMatrix, b::PVector) + hypre_x = HYPREVector(x) + HYPRE.solve!(solver, hypre_x, HYPREMatrix(A), HYPREVector(b)) + copy!(x, hypre_x) + return x +end + +end # module HYPREPartitionedArrays diff --git a/src/HYPRE.jl b/src/HYPRE.jl index ebffd11..2c2a3e8 100644 --- a/src/HYPRE.jl +++ b/src/HYPRE.jl @@ -3,9 +3,6 @@ module HYPRE using MPI: MPI -using PartitionedArrays: PartitionedArrays, AbstractLocalIndices, MPIArray, PSparseMatrix, - PVector, SplitMatrix, ghost_to_global, local_values, own_to_global, own_values, - partition using SparseArrays: SparseArrays, SparseMatrixCSC, nonzeros, nzrange, rowvals using SparseMatricesCSR: SparseMatrixCSR, colvals @@ -331,299 +328,6 @@ function Base.copy!(dst::HYPREVector, src::Vector{HYPRE_Complex}) return dst end -################################################## -# PartitionedArrays.PSparseMatrix -> HYPREMatrix # -################################################## - -function subarray_unsafe_supported() - # Wrapping of SubArrays as raw pointers may or may not be supported - # depending on the Julia version. If this is not supported, we have to fall - # back to allocation of an intermediate buffer. This logic can be removed if - # HYPRE.jl drops support for Julia < 1.9. - return VERSION >= v"1.9.0" -end - -function Internals.to_hypre_data( - A::SplitMatrix{<:SparseMatrixCSC}, r::AbstractLocalIndices, c::AbstractLocalIndices - ) - # Own/ghost to global index mappings - own_to_global_row = own_to_global(r) - own_to_global_col = own_to_global(c) - ghost_to_global_col = ghost_to_global(c) - - # HYPRE requires contiguous row indices - ilower = own_to_global_row[1] - iupper = own_to_global_row[end] - @assert iupper - ilower + 1 == length(own_to_global_row) - - # Extract sparse matrices from the SplitMatrix. We are only interested in the owned - # rows, so only consider own-own and own-ghost blocks. - Aoo = A.blocks.own_own::SparseMatrixCSC - Aoo_rows = rowvals(Aoo) - Aoo_vals = nonzeros(Aoo) - Aog = A.blocks.own_ghost::SparseMatrixCSC - Aog_rows = rowvals(Aog) - Aog_vals = nonzeros(Aog) - @assert size(Aoo, 1) == size(Aog, 1) == length(own_to_global_row) - - # Initialize the data buffers HYPRE wants - nrows = HYPRE_Int(length(own_to_global_row)) # Total number of rows - ncols = zeros(HYPRE_Int, nrows) # Number of colums for each row - rows = collect(HYPRE_BigInt, ilower:iupper) # The row indices - # cols = Vector{HYPRE_BigInt}(undef, nnz) # The column indices - # values = Vector{HYPRE_Complex}(undef, nnz) # The values - - # First pass to count nnz per row (note that global column indices and column - # permutation doesn't matter for this pass) - @inbounds for own_col in 1:size(Aoo, 2) - for k in nzrange(Aoo, own_col) - own_row = Aoo_rows[k] - ncols[own_row] += 1 - end - end - @inbounds for ghost_col in 1:size(Aog, 2) - for k in nzrange(Aog, ghost_col) - own_row = Aog_rows[k] - ncols[own_row] += 1 - end - end - - # Initialize remaining buffers now that nnz is known - nnz = sum(ncols) - cols = Vector{HYPRE_BigInt}(undef, nnz) - values = Vector{HYPRE_Complex}(undef, nnz) - - # Keep track of the last index used for every row - lastinds = zeros(Int, nrows) - cumsum!((@view lastinds[2:end]), (@view ncols[1:end-1])) - - # Second pass to populate the output. Here we need to map column - # indices from own/ghost to global - @inbounds for own_col in 1:size(Aoo, 2) - for k in nzrange(Aoo, own_col) - own_row = Aoo_rows[k] - i = lastinds[own_row] += 1 - values[i] = Aoo_vals[k] - cols[i] = own_to_global_col[own_col] - end - end - @inbounds for ghost_col in 1:size(Aog, 2) - for k in nzrange(Aog, ghost_col) - own_row = Aog_rows[k] - i = lastinds[own_row] += 1 - values[i] = Aog_vals[k] - cols[i] = ghost_to_global_col[ghost_col] - end - end - - # Sanity checks and return - @assert nrows == length(ncols) == length(rows) - return nrows, ncols, rows, cols, values -end - -function Internals.to_hypre_data( - A::SplitMatrix{<:SparseMatrixCSR}, r::AbstractLocalIndices, c::AbstractLocalIndices - ) - # Own/ghost to global index mappings - own_to_global_row = own_to_global(r) - own_to_global_col = own_to_global(c) - ghost_to_global_col = ghost_to_global(c) - - # HYPRE requires contiguous row indices - ilower = own_to_global_row[1] - iupper = own_to_global_row[end] - @assert iupper - ilower + 1 == length(own_to_global_row) - - # Extract sparse matrices from the SplitMatrix. We are only interested in the owned - # rows, so only consider own-own and own-ghost blocks. - Aoo = A.blocks.own_own::SparseMatrixCSR - Aoo_cols = colvals(Aoo) - Aoo_vals = nonzeros(Aoo) - Aog = A.blocks.own_ghost::SparseMatrixCSR - Aog_cols = colvals(Aog) - Aog_vals = nonzeros(Aog) - @assert size(Aoo, 1) == size(Aog, 1) == length(own_to_global_row) - - # Initialize the data buffers HYPRE wants - nnz = SparseArrays.nnz(Aoo) + SparseArrays.nnz(Aog) - nrows = HYPRE_Int(iupper - ilower + 1) # Total number of rows - ncols = zeros(HYPRE_Int, nrows) # Number of columns for each row - rows = collect(HYPRE_BigInt, ilower:iupper) # The row indices - cols = Vector{HYPRE_BigInt}(undef, nnz) # The column indices - values = Vector{HYPRE_Complex}(undef, nnz) # The values - - # For CSR we only need a single pass to over the owned rows to collect everything - i = 0 - for own_row in 1:size(Aoo, 1) - nzro = nzrange(Aoo, own_row) - nzrg = nzrange(Aog, own_row) - ncols[own_row] = length(nzro) + length(nzrg) - for k in nzro - i += 1 - own_col = Aoo_cols[k] - cols[i] = own_to_global_col[own_col] - values[i] = Aoo_vals[k] - end - for k in nzrg - i += 1 - ghost_col = Aog_cols[k] - cols[i] = ghost_to_global_col[ghost_col] - values[i] = Aog_vals[k] - end - end - - # Sanity checks and return - @assert nnz == i - @assert nrows == length(ncols) == length(rows) - return nrows, ncols, rows, cols, values -end - -function Internals.get_comm(A::Union{PSparseMatrix{<:Any,<:M}, PVector{<:Any,<:M}}) where M <: MPIArray - return partition(A).comm -end - -Internals.get_comm(_::Union{PSparseMatrix,PVector}) = MPI.COMM_SELF - -function Internals.get_proc_rows(A::Union{PSparseMatrix, PVector}) - if A isa PVector - r = A.index_partition - else - r = A.row_partition - end - ilower::HYPRE_BigInt = typemax(HYPRE_BigInt) - iupper::HYPRE_BigInt = typemin(HYPRE_BigInt) - map(r) do a - # This is a map over the local process' owned indices. For MPI it will - # be a single value but for DebugArray / Array it will have multiple - # values. - o_to_g = own_to_global(a) - ilower_part = o_to_g[1] - iupper_part = o_to_g[end] - ilower = min(ilower, convert(HYPRE_BigInt, ilower_part)) - iupper = max(iupper, convert(HYPRE_BigInt, iupper_part)) - end - return ilower, iupper -end - -function HYPREMatrix(B::PSparseMatrix) - # Use the same communicator as the matrix - comm = Internals.get_comm(B) - # Fetch rows owned by this process - ilower, iupper = Internals.get_proc_rows(B) - # Create the IJ matrix - A = HYPREMatrix(comm, ilower, iupper) - # Set all the values - map(local_values(B), B.row_partition, B.col_partition) do Bv, Br, Bc - nrows, ncols, rows, cols, values = Internals.to_hypre_data(Bv, Br, Bc) - @check HYPRE_IJMatrixSetValues(A, nrows, ncols, rows, cols, values) - return nothing - end - # Finalize - Internals.assemble_matrix(A) - return A -end - -############################################ -# PartitionedArrays.PVector -> HYPREVector # -############################################ - -function HYPREVector(v::PVector) - # Use the same communicator as the matrix - comm = Internals.get_comm(v) - # Fetch rows owned by this process - ilower, iupper = Internals.get_proc_rows(v) - # Create the IJ vector - b = HYPREVector(comm, ilower, iupper) - # Set all the values - map(own_values(v), v.index_partition) do vo, vr - o_to_g = own_to_global(vr) - - ilower_part = o_to_g[1] - iupper_part = o_to_g[end] - - # Option 1: Set all values - nvalues = HYPRE_Int(iupper_part - ilower_part + 1) - indices = collect(HYPRE_BigInt, ilower_part:iupper_part) - # TODO: Could probably just pass the full vector even if it is too long - # values = convert(Vector{HYPRE_Complex}, vv) - values = collect(HYPRE_Complex, vo) - - # # Option 2: Set only non-zeros - # indices = HYPRE_BigInt[] - # values = HYPRE_Complex[] - # for (i, vi) in zip(ilower_part:iupper_part, vo) - # if !iszero(vi) - # push!(indices, i) - # push!(values, vi) - # end - # end - # nvalues = length(indices) - - @check HYPRE_IJVectorSetValues(b, nvalues, indices, values) - return nothing - end - # Finalize - Internals.assemble_vector(b) - return b -end - -function Internals.copy_check(dst::HYPREVector, src::PVector) - il_dst, iu_dst = Internals.get_proc_rows(dst) - il_src, iu_src = Internals.get_proc_rows(src) - if il_dst != il_src && iu_dst != iu_src - # TODO: Why require this? - throw(ArgumentError( - "row owner mismatch between dst ($(il_dst:iu_dst)) and src ($(il_src:iu_src))" - )) - end -end - -# TODO: Other eltypes could be support by using a intermediate buffer -function Base.copy!(dst::PVector{<:AbstractVector{HYPRE_Complex}}, src::HYPREVector) - Internals.copy_check(src, dst) - map(own_values(dst), dst.index_partition) do ov, vr - o_to_g = own_to_global(vr) - il_src_part = o_to_g[1] - iu_src_part = o_to_g[end] - nvalues = HYPRE_Int(iu_src_part - il_src_part + 1) - indices = collect(HYPRE_BigInt, il_src_part:iu_src_part) - if subarray_unsafe_supported() - values = ov - else - values = collect(HYPRE_Complex, ov) - end - @check HYPRE_IJVectorGetValues(src, nvalues, indices, values) - if !subarray_unsafe_supported() - @. ov = values - end - end - return dst -end - -function Base.copy!(dst::HYPREVector, src::PVector{<:AbstractVector{HYPRE_Complex}}) - Internals.copy_check(dst, src) - # Re-initialize the vector - @check HYPRE_IJVectorInitialize(dst) - map(own_values(src), src.index_partition) do ov, vr - o_to_g = own_to_global(vr) - ilower_src_part = o_to_g[1] - iupper_src_part = o_to_g[end] - nvalues = HYPRE_Int(iupper_src_part - ilower_src_part + 1) - indices = collect(HYPRE_BigInt, ilower_src_part:iupper_src_part) - if subarray_unsafe_supported() - values = ov - else - values = collect(HYPRE_Complex, ov) - end - @check HYPRE_IJVectorSetValues(dst, nvalues, indices, values) - end - # TODO: It shouldn't be necessary to assemble here since we only set owned rows (?) - # @check HYPRE_IJVectorAssemble(dst) - # TODO: Necessary to recreate the ParVector? Running some examples it seems like it is - # not needed. - return dst -end - #################### ## HYPREAssembler ## @@ -797,4 +501,9 @@ end include("solvers.jl") include("solver_options.jl") +# Compatibility for Julias that doesn't support package extensions +if !(isdefined(Base, :get_extension)) + include("../ext/HYPREPartitionedArrays.jl") +end + end # module HYPRE diff --git a/src/solvers.jl b/src/solvers.jl index 9c0659b..1c8df03 100644 --- a/src/solvers.jl +++ b/src/solvers.jl @@ -51,24 +51,6 @@ See also [`solve`](@ref). solve!(pcg::HYPRESolver, x::HYPREVector, A::HYPREMatrix, ::HYPREVector) -###################################### -# PartitionedArrays solver interface # -###################################### - -# TODO: Would it be useful with a method that copied the solution to b instead? - -function solve(solver::HYPRESolver, A::PSparseMatrix, b::PVector) - hypre_x = solve(solver, HYPREMatrix(A), HYPREVector(b)) - x = copy!(similar(b, HYPRE_Complex), hypre_x) - return x -end -function solve!(solver::HYPRESolver, x::PVector, A::PSparseMatrix, b::PVector) - hypre_x = HYPREVector(x) - solve!(solver, hypre_x, HYPREMatrix(A), HYPREVector(b)) - copy!(x, hypre_x) - return x -end - ######################################## # SparseMatrixCS(C|R) solver interface # ########################################