Merge pull request #82 from jipolanco/pkgextensions

Use package extensions, require Julia 1.9

.github/workflows/ci.yml

@@ -24,7 +24,6 @@ jobs:
       matrix:
         experimental: [false]
         version:
-          - '1.8'
           - '1.9'
         os:
           - ubuntu-latest

Project.toml

@@ -1,7 +1,7 @@
 name = "PencilArrays"
 uuid = "0e08944d-e94e-41b1-9406-dcf66b6a9d2e"
-authors = ["Juan Ignacio Polanco <[email protected]> and contributors"]
-version = "0.18.0"
+authors = ["Juan Ignacio Polanco <[email protected]> and contributors"]
+version = "0.18.1"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
@@ -11,25 +11,33 @@ MPI = "da04e1cc-30fd-572f-bb4f-1f8673147195"
 OffsetArrays = "6fe1bfb0-de20-5000-8ca7-80f57d26f881"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
-Requires = "ae029012-a4dd-5104-9daa-d747884805df"
 StaticArrayInterface = "0d7ed370-da01-4f52-bd93-41d350b8b718"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 StaticPermutations = "15972242-4b8f-49a0-b8a1-9ac0e7a1a45d"
 Strided = "5e0ebb24-38b0-5f93-81fe-25c709ecae67"
 TimerOutputs = "a759f4b9-e2f1-59dc-863e-4aeb61b1ea8f"
 VersionParsing = "81def892-9a0e-5fdd-b105-ffc91e053289"
 
+[weakdeps]
+DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e"
+HDF5 = "f67ccb44-e63f-5c2f-98bd-6dc0ccc4ba2f"
+
+[extensions]
+PencilArraysDiffEqExt = ["DiffEqBase"]
+PencilArraysHDF5Ext = ["HDF5"]
+
 [compat]
 Adapt = "3"
+DiffEqBase = "6"
+HDF5 = "0.16"
 JSON3 = "1.4"
 MPI = "0.20"
 OffsetArrays = "1"
 Reexport = "1"
-Requires = "1"
 StaticArrayInterface = "1"
 StaticArrays = "1"
 StaticPermutations = "0.3"
 Strided = "2"
 TimerOutputs = "0.5"
 VersionParsing = "1"
-julia = "1.8"
+julia = "1.9"

ext/PencilArraysDiffEqExt.jl

@@ -0,0 +1,11 @@
+module PencilArraysDiffEqExt
+
+using PencilArrays: PencilArray, length_global
+using DiffEqBase
+
+# This is used for adaptive timestepping in DifferentialEquations.jl.
+# Without this, each MPI process may choose a different dt, leading to
+# catastrophic consequences!
+DiffEqBase.recursive_length(u::PencilArray) = length_global(u)
+
+end

ext/PencilArraysHDF5Ext.jl

@@ -0,0 +1,258 @@
+module PencilArraysHDF5Ext
+
+using MPI
+using HDF5
+using PencilArrays
+using PencilArrays.PencilIO
+using PencilArrays: MaybePencilArrayCollection, collection_size
+using StaticArrays: SVector
+using TimerOutputs
+
+function PencilIO.PHDF5Driver(;
+        fcpl = HDF5.FileCreateProperties(),
+        fapl = HDF5.FileAccessProperties(),
+    )
+    # We set fapl.fclose_degree if it hasn't been explicitly set.
+    if !_is_set(fapl, Val(:fclose_degree))
+        # This is the default in HDF5.jl -- makes sense due to GC.
+        fapl.fclose_degree = :strong
+    end
+    PencilIO._PHDF5Driver(fcpl, fapl)
+end
+
+# TODO Is there a better way to check if fapl.fclose_degree has already
+# been set??
+function _is_set(fapl::HDF5.FileAccessProperties, ::Val{:fclose_degree})
+    id = fapl.id
+    degree = Ref{Cint}()
+    status = ccall(
+        (:H5Pget_fclose_degree, HDF5.API.libhdf5), HDF5.API.herr_t,
+        (HDF5.API.hid_t, Ref{Cint}), id, degree)
+    # A negative value means failure, which we interpret here as meaning that
+    # "fclose_degree" has not been set.
+    status ≥ 0
+end
+
+PencilIO.hdf5_has_parallel() = HDF5.has_parallel()
+
+function keywords_to_h5open(; kws...)
+    flags, other_kws = PencilIO.keywords_to_open(; kws...)
+    (
+        flags.read,
+        flags.write,
+        flags.create,
+        flags.truncate,
+        flags.append,
+    ), other_kws
+end
+
+function Base.open(D::PHDF5Driver, filename::AbstractString, comm::MPI.Comm; kw...)
+    mode_args, other_kws = keywords_to_h5open(; kw...)
+    info = MPI.Info(other_kws...)
+    fcpl = D.fcpl
+    fapl = D.fapl
+    mpio = HDF5.Drivers.MPIO(comm, info)
+    HDF5.Drivers.set_driver!(fapl, mpio)  # fails if no parallel support
+    swmr = false
+
+    # The code below is adapted from h5open in HDF5.jl v0.15
+    # TODO propose alternative h5open for HDF5.jl, taking keyword arguments `read`, `write`, ...
+    # Then we wouldn't need to copy code from HDF5.jl...
+    rd, wr, cr, tr, ff = mode_args
+    if ff && !wr
+        error("HDF5 does not support appending without writing")
+    end
+
+    fid = if cr && (tr || !isfile(filename))
+        flag = swmr ? HDF5.API.H5F_ACC_TRUNC | HDF5.API.H5F_ACC_SWMR_WRITE :
+                      HDF5.API.H5F_ACC_TRUNC
+        HDF5.API.h5f_create(filename, flag, fcpl, fapl)
+    else
+        HDF5.ishdf5(filename) ||
+            error("unable to determine if $filename is accessible in the HDF5 format (file may not exist)")
+        flag = if wr
+            swmr ? HDF5.API.H5F_ACC_RDWR | HDF5.API.H5F_ACC_SWMR_WRITE :
+                   HDF5.API.H5F_ACC_RDWR
+        else
+            swmr ? HDF5.API.H5F_ACC_RDONLY | HDF5.API.H5F_ACC_SWMR_READ :
+                   HDF5.API.H5F_ACC_RDONLY
+        end
+        HDF5.API.h5f_open(filename, flag, fapl)
+    end
+
+    close(fapl)
+    close(fcpl)
+
+    HDF5.File(fid, filename)
+end
+
+function Base.setindex!(
+        g::Union{HDF5.File, HDF5.Group}, x::MaybePencilArrayCollection,
+        name::AbstractString;
+        chunks=false, collective=true, prop_pairs...,
+    )
+    to = timer(pencil(x))
+
+    @timeit_debug to "Write HDF5" begin
+
+    check_phdf5_file(g, x)
+
+    # Add extra property lists if required by keyword args.
+    # TODO avoid using Dict?
+    props = Dict{Symbol,Any}(pairs(prop_pairs))
+
+    if chunks && !haskey(prop_pairs, :chunk)
+        chunk = h5_chunk_size(x, MemoryOrder())
+        props[:chunk] = chunk
+    end
+
+    if collective && !haskey(prop_pairs, :dxpl_mpio)
+        props[:dxpl_mpio] = :collective
+    end
+
+    dims_global = h5_dataspace_dims(x)
+    @timeit_debug to "create dataset" dset =
+        create_dataset(g, name, h5_datatype(x), dataspace(dims_global); props...)
+    inds = range_local(x, MemoryOrder())
+    @timeit_debug to "write data" to_hdf5(dset, x, inds)
+    @timeit_debug to "write metadata" write_metadata(dset, x)
+
+    end
+
+    x
+end
+
+# Write metadata as HDF5 attributes attached to a dataset.
+# Note that this is a collective operation (all processes must call this).
+function write_metadata(dset::HDF5.Dataset, x)
+    meta = PencilIO.metadata(x)
+    for (name, val) in pairs(meta)
+        dset[string(name)] = to_hdf5(val)
+    end
+    dset
+end
+
+to_hdf5(val) = val
+to_hdf5(val::Tuple{}) = false  # empty tuple
+to_hdf5(val::Tuple) = SVector(val)
+to_hdf5(::Nothing) = false
+
+function Base.read!(g::Union{HDF5.File, HDF5.Group}, x::MaybePencilArrayCollection,
+                    name::AbstractString; collective=true, prop_pairs...)
+    to = timer(pencil(x))
+
+    @timeit_debug to "Read HDF5" begin
+
+    dapl = HDF5.DatasetAccessProperties(; prop_pairs...)
+    dxpl = HDF5.DatasetTransferProperties(; prop_pairs...)
+
+    # Add extra property lists if required by keyword args.
+    if collective && !haskey(prop_pairs, :dxpl_mpio)
+        dxpl.dxpl_mpio = :collective
+    end
+
+    dims_global = h5_dataspace_dims(x)
+    @timeit_debug to "open dataset" dset = open_dataset(g, string(name), dapl, dxpl)
+    check_phdf5_file(parent(dset), x)
+
+    if dims_global != size(dset)
+        throw(DimensionMismatch(
+            "incompatible dimensions of HDF5 dataset and PencilArray"))
+    end
+
+    inds = range_local(x, MemoryOrder())
+    @timeit_debug to "read data" from_hdf5!(dset, x, inds)
+
+    end
+
+    x
+end
+
+function check_phdf5_file(g, x)
+    fapl = HDF5.get_access_properties(HDF5.file(g))
+    driver = HDF5.Drivers.get_driver(fapl)
+    if driver isa HDF5.Drivers.MPIO
+        comm = driver.comm
+        if MPI.Comm_compare(comm, get_comm(x)) ∉ (MPI.IDENT, MPI.CONGRUENT)
+            throw(ArgumentError(
+                "incompatible MPI communicators of HDF5 file and PencilArray"
+            ))
+        end
+    else
+        error("HDF5 file was not opened with the MPIO driver")
+    end
+    close(fapl)
+    nothing
+end
+
+to_hdf5(dset, x::PencilArray, inds) = dset[inds...] = parent(x)
+
+function from_hdf5!(dset, x::PencilArray, inds)
+    u = parent(x)
+
+    if stride(u, 1) != 1
+        u .= dset[inds...]  # short and easy version (but allocates!)
+        return x
+    end
+
+    # The following is adapted from one of the _getindex() in HDF5.jl.
+    HDF5Scalar = HDF5.ScalarType
+    T = eltype(x)
+    if !(T <: Union{HDF5Scalar, Complex{<:HDF5Scalar}})
+        error("Dataset indexing (hyperslab) is available only for bits types")
+    end
+
+    dsel_id = HDF5.hyperslab(dset, inds...)
+    memtype = HDF5.datatype(u)
+    memspace = HDF5.dataspace(u)
+
+    try
+        # This only works for stride-1 arrays.
+        HDF5.API.h5d_read(dset.id, memtype.id, memspace.id, dsel_id, dset.xfer, u)
+    finally
+        close(memtype)
+        close(memspace)
+        HDF5.API.h5s_close(dsel_id)
+    end
+
+    x
+end
+
+# Define variants for collections.
+for func in (:from_hdf5!, :to_hdf5)
+    @eval function $func(dset, col::PencilArrayCollection, inds_in)
+        for I in CartesianIndices(collection_size(col))
+            inds = (inds_in..., Tuple(I)...)
+            $func(dset, col[I], inds)
+        end
+    end
+end
+
+h5_datatype(x::PencilArray) = datatype(eltype(x))
+h5_datatype(x::PencilArrayCollection) = h5_datatype(first(x))
+
+h5_dataspace_dims(x::PencilArray) = size_global(x, MemoryOrder())
+h5_dataspace_dims(x::PencilArrayCollection) =
+    (h5_dataspace_dims(first(x))..., collection_size(x)...)
+
+function h5_chunk_size(x::PencilArray, order = MemoryOrder())
+    # Determine chunk size for writing to HDF5 dataset.
+    # The idea is that each process writes to a single separate chunk of the
+    # dataset, of size `dims_local`.
+    # This only works if the data is ideally balanced among processes, i.e. if
+    # the local dimensions of the dataset are the same for all processes.
+    dims_local = size_local(x, order)
+
+    # In the general case that the data is not well balanced, we take the
+    # minimum size along each dimension.
+    chunk = MPI.Allreduce(collect(dims_local), min, get_comm(x))
+
+    N = ndims(x)
+    @assert length(chunk) == N
+    ntuple(d -> chunk[d], Val(N))
+end
+
+h5_chunk_size(x::PencilArrayCollection, args...) =
+    (h5_chunk_size(first(x), args...)..., collection_size(x)...)
+
+end

src/PencilArrays.jl

@@ -5,7 +5,6 @@ using OffsetArrays
 using Reexport
 using StaticPermutations
 using TimerOutputs
-using Requires: @require
 
 using Base: @propagate_inbounds
 import Adapt
@@ -57,13 +56,4 @@ include("Transpositions/Transpositions.jl")
 
 include("PencilIO/PencilIO.jl")
 
-function __init__()
-    @require DiffEqBase="2b5f629d-d688-5b77-993f-72d75c75574e" @eval begin
-        # This is used for adaptive timestepping in DifferentialEquations.jl.
-        # Without this, each MPI process may choose a different dt, leading to
-        # catastrophic consequences!
-        DiffEqBase.recursive_length(u::PencilArray) = length_global(u)
-    end
-end
-
 end

src/PencilIO/PencilIO.jl

@@ -4,7 +4,6 @@ using ..PencilArrays
 import ..PencilArrays: MaybePencilArrayCollection, collection_size, collection
 
 using MPI
-using Requires: @require
 using StaticArrays: SVector
 using TimerOutputs
 
@@ -73,9 +72,6 @@ function keywords_to_open(; read=nothing, write=nothing, create=nothing,
 end
 
 include("mpi_io.jl")
-
-function __init__()
-    @require HDF5="f67ccb44-e63f-5c2f-98bd-6dc0ccc4ba2f" include("hdf5.jl")
-end
+include("hdf5.jl")  # actual implementation is in ../../ext/PencilArraysHDF5Ext.jl
 
 end