HDF5Interface.h

// Copyright (C) 2012 Chris N. Richardson and Garth N. Wells
//
// This file is part of DOLFINx (https://www.fenicsproject.org)
//
// SPDX-License-Identifier:    LGPL-3.0-or-later
 
#pragma once
 
#include <array>
#include <cassert>
#include <chrono>
#include <cstdint>
#include <dolfinx/common/log.h>
#include <filesystem>
#include <hdf5.h>
#include <mpi.h>
#include <numeric>
#include <string>
#include <vector>
 
namespace dolfinx::io::hdf5
{
template <typename T>
hid_t hdf5_type()
{
  if constexpr (std::is_same_v<T, float>)
    return H5T_NATIVE_FLOAT;
  else if constexpr (std::is_same_v<T, double>)
    return H5T_NATIVE_DOUBLE;
  else if constexpr (std::is_same_v<T, std::int32_t>)
    return H5T_NATIVE_INT32;
  else if constexpr (std::is_same_v<T, std::uint32_t>)
    return H5T_NATIVE_UINT32;
  else if constexpr (std::is_same_v<T, std::int64_t>)
    return H5T_NATIVE_INT64;
  else if constexpr (std::is_same_v<T, std::uint64_t>)
    return H5T_NATIVE_UINT64;
  else if constexpr (std::is_same_v<T, std::uint8_t>)
    return H5T_NATIVE_UINT8;
  else if constexpr (std::is_same_v<T, std::size_t>)
  {
    throw std::runtime_error(
        "Cannot determine size of std::size_t. std::size_t is not the same "
        "size as long or int.");
  }
  else
  {
    throw std::runtime_error("Cannot get HDF5 primitive data type. No "
                             "specialised function for this data type.");
  }
}

hid_t open_file(MPI_Comm comm, const std::filesystem::path& filename,
                const std::string& mode, bool use_mpi_io);

void close_file(hid_t handle);

void flush_file(hid_t handle);

std::filesystem::path get_filename(hid_t handle);

bool has_dataset(hid_t handle, const std::string& dataset_path);

void set_attribute(hid_t handle, const std::string& attr_name,
                   const std::string& value);
void set_attribute(hid_t handle, const std::string& attr_name,
                   const std::vector<std::int32_t>& value);
void set_attribute(hid_t handle, const std::string& attr_name,
                   std::int32_t value);

hid_t open_dataset(hid_t handle, const std::string& path);

std::vector<std::int64_t> get_dataset_shape(hid_t handle,
                                            const std::string& dataset_path);

void set_mpi_atomicity(hid_t handle, bool atomic);

bool get_mpi_atomicity(hid_t handle);

void add_group(hid_t handle, const std::string& dataset_path);

template <typename T>
void write_dataset(hid_t file_handle, const std::string& dataset_path,
                   const T* data, std::array<std::int64_t, 2> range,
                   const std::vector<int64_t>& global_size, bool use_mpi_io,
                   bool use_chunking)
{
  // Check that group exists and recursively create if required
  const std::string group_name(dataset_path, 0, dataset_path.rfind('/'));
  add_group(file_handle, group_name);
 
  // Data rank
  const int rank = global_size.size();
  assert(rank != 0);
  if (rank > 2)
  {
    throw std::runtime_error("Cannot write dataset to HDF5 file"
                             "Only rank 1 and rank 2 dataset are supported");
  }
 
  // Get HDF5 data type
  const hid_t h5type = hdf5::hdf5_type<T>();
 
  // Hyperslab selection parameters
  std::vector<hsize_t> count(global_size.begin(), global_size.end());
  count[0] = range[1] - range[0];
 
  // Data offsets
  std::vector<hsize_t> offset(rank, 0);
  offset[0] = range[0];
 
  // Dataset dimensions
  const std::vector<hsize_t> dimsf(global_size.begin(), global_size.end());
 
  // Writing into an existing dataset
  if (has_dataset(file_handle, dataset_path))
  {
    // Resize to new global size
    hid_t dset_id = hdf5::open_dataset(file_handle, dataset_path.c_str());
    H5Dset_extent(dset_id, dimsf.data());
    H5Dclose(dset_id);
  }
  else
  {
    std::vector<hsize_t> maxdims(dimsf.begin(), dimsf.end());
 
    // Set chunking parameters
    hid_t chunking_properties;
    if (use_chunking)
    {
      // Make array extensible, if chunking is set
      std::fill(maxdims.begin(), maxdims.end(), H5S_UNLIMITED);
 
      // Set chunk size and limit to 1kB min/1MB max
      hsize_t chunk_size = dimsf[0] / 2;
      if (chunk_size > 1048576)
        chunk_size = 1048576;
      if (chunk_size < 1024)
        chunk_size = 1024;
 
      hsize_t chunk_dims[2] = {chunk_size, dimsf[1]};
      chunking_properties = H5Pcreate(H5P_DATASET_CREATE);
      H5Pset_chunk(chunking_properties, rank, chunk_dims);
    }
    else
      chunking_properties = H5P_DEFAULT;
 
    // Create a global data space
    const hid_t filespace0
        = H5Screate_simple(rank, dimsf.data(), maxdims.data());
    if (filespace0 == H5I_INVALID_HID)
      throw std::runtime_error("Failed to create HDF5 data space");
 
    // Create global dataset (using dataset_path)
    const hid_t dset_id
        = H5Dcreate2(file_handle, dataset_path.c_str(), h5type, filespace0,
                     H5P_DEFAULT, chunking_properties, H5P_DEFAULT);
    if (dset_id == H5I_INVALID_HID)
      throw std::runtime_error("Failed to create HDF5 global dataset.");
 
    // Close global data space
    if (H5Sclose(filespace0) < 0)
      throw std::runtime_error("Failed to close HDF5 global data space.");
 
    if (use_chunking)
    {
      // Close chunking properties
      if (H5Pclose(chunking_properties) < 0)
        throw std::runtime_error("Failed to close HDF5 chunking properties.");
    }
 
    // Close dataset collectively
    if (H5Dclose(dset_id) < 0)
      throw std::runtime_error("Failed to close HDF5 dataset.");
  }
 
  // Reopen dataset and write data
  hid_t dset_id = hdf5::open_dataset(file_handle, dataset_path.c_str());
  assert(dset_id != H5I_INVALID_HID);
 
  hid_t dataspace = H5Dget_space(dset_id);
  if (dataspace == H5I_INVALID_HID)
    throw std::runtime_error("Failed to open HDF5 data space.");
 
  herr_t status = H5Sselect_hyperslab(dataspace, H5S_SELECT_SET, offset.data(),
                                      nullptr, count.data(), nullptr);
  if (status < 0)
    throw std::runtime_error("Failed to create HDF5 dataspace.");
 
  // Set parallel access
  const hid_t plist_id = H5Pcreate(H5P_DATASET_XFER);
  if (use_mpi_io)
  {
    if (herr_t status = H5Pset_dxpl_mpio(plist_id, H5FD_MPIO_COLLECTIVE);
        status < 0)
    {
      throw std::runtime_error(
          "Failed to set HDF5 data transfer property list.");
    }
  }
 
  // Create a local data space
  const hid_t memspace = H5Screate_simple(rank, count.data(), nullptr);
  if (memspace == H5I_INVALID_HID)
    throw std::runtime_error("Failed to create HDF5 local data space.");
 
  // Write local dataset
  if (H5Dwrite(dset_id, h5type, memspace, dataspace, plist_id, data) < 0)
  {
    throw std::runtime_error("Failed to write HDF5 local dataset.");
  }
 
  H5Sclose(memspace);
  H5Sclose(dataspace);
 
  H5Dclose(dset_id);
  // Release file-access template
  if (H5Pclose(plist_id) < 0)
    throw std::runtime_error("Failed to release HDF5 file-access template.");
}

template <typename T>
std::vector<T> read_dataset(hid_t dset_id, std::array<std::int64_t, 2> range,
                            bool allow_cast)
{
  auto timer_start = std::chrono::system_clock::now();
 
  if (!allow_cast)
  {
    // Check that HDF5 dataset type and the type T are the same
 
    hid_t dtype = H5Dget_type(dset_id);
    if (dtype == H5I_INVALID_HID)
      throw std::runtime_error("Failed to get HDF5 data type.");
    if (htri_t eq = H5Tequal(dtype, hdf5::hdf5_type<T>()); eq < 0)
      throw std::runtime_error("HDF5 datatype equality test failed.");
    else if (!eq)
    {
      H5Tclose(dtype);
      throw std::runtime_error("Wrong type for reading from HDF5. Use \"h5ls "
                               "-v\" to inspect the types in the HDF5 file.");
    }
  }
 
  // Open dataspace
  hid_t dataspace = H5Dget_space(dset_id);
  if (dataspace == H5I_INVALID_HID)
    throw std::runtime_error("Failed to open HDF5 data space.");
 
  // Get rank of data set
  int rank = H5Sget_simple_extent_ndims(dataspace);
  if (rank < 1)
    throw std::runtime_error("Failed to get rank of data space.");
  else if (rank > 2)
    spdlog::warn("io::hdf5::read_dataset untested for rank > 2.");
 
  // Allocate data for shape
  std::vector<hsize_t> shape(rank);
 
  // Get size in each dimension
  if (int ndims = H5Sget_simple_extent_dims(dataspace, shape.data(), nullptr);
      ndims != rank)
  {
    throw std::runtime_error("Failed to get dimensionality of dataspace.");
  }
 
  // Hyperslab selection
  std::vector<hsize_t> offset(rank, 0);
  std::vector<hsize_t> count = shape;
  if (range[0] != -1 and range[1] != -1)
  {
    offset[0] = range[0];
    count[0] = range[1] - range[0];
  }
  else
    offset[0] = 0;
 
  // Select a block in the dataset beginning at offset[], with
  // size=count[]
  if (herr_t status
      = H5Sselect_hyperslab(dataspace, H5S_SELECT_SET, offset.data(), nullptr,
                            count.data(), nullptr);
      status < 0)
  {
    throw std::runtime_error("Failed to select HDF5 hyperslab.");
  }
 
  // Create a memory dataspace
  hid_t memspace = H5Screate_simple(rank, count.data(), nullptr);
  if (memspace == H5I_INVALID_HID)
    throw std::runtime_error("Failed to create HDF5 dataspace.");
 
  // Create local data to read into
  std::vector<T> data(
      std::reduce(count.begin(), count.end(), 1, std::multiplies{}));
 
  // Read data on each process
  hid_t h5type = hdf5::hdf5_type<T>();
  if (herr_t status
      = H5Dread(dset_id, h5type, memspace, dataspace, H5P_DEFAULT, data.data());
      status < 0)
  {
    throw std::runtime_error("Failed to read HDF5 data.");
  }
 
  // Close dataspace
  if (herr_t status = H5Sclose(dataspace); status < 0)
    throw std::runtime_error("Failed to close HDF5 dataspace.");
 
  // Close memspace
  if (herr_t status = H5Sclose(memspace); status < 0)
    throw std::runtime_error("Failed to close HDF5 memory space.");
 
  auto timer_end = std::chrono::system_clock::now();
  std::chrono::duration<double> dt = (timer_end - timer_start);
  double data_rate = data.size() * sizeof(T) / (1e6 * dt.count());
  spdlog::info("HDF5 Read data rate: {} MB/s", data_rate);
 
  return data;
}
} // namespace dolfinx::io::hdf5