options.hpp

Go to the documentation of this file.

 
#pragma once
 
#include <core.hpp>
#include <utils/fqn.hpp>
#include <utils/getenv.hpp>
 
#include <torch/torch.h>
 
namespace iganet {
 
struct half {};
 
template <typename T> inline constexpr torch::Dtype dtype();
 
template <> inline constexpr torch::Dtype dtype<bool>() { return torch::kBool; }
 
template <> inline constexpr torch::Dtype dtype<char>() { return torch::kChar; }
 
template <> inline constexpr torch::Dtype dtype<short>() {
  return torch::kShort;
}
 
template <> inline constexpr torch::Dtype dtype<int>() { return torch::kInt; }
 
  template <> inline constexpr torch::Dtype dtype<long>() { return torch::kLong; }
 
  template <> inline constexpr torch::Dtype dtype<long long>() { return torch::kLong; }
 
template <> inline constexpr torch::Dtype dtype<half>() { return torch::kHalf; }
 
template <> inline constexpr torch::Dtype dtype<float>() {
  return torch::kFloat;
}
 
template <> inline constexpr torch::Dtype dtype<double>() {
  return torch::kDouble;
}
 
template <> inline constexpr torch::Dtype dtype<std::complex<half>>() {
  return at::kComplexHalf;
}
 
template <> inline constexpr torch::Dtype dtype<std::complex<float>>() {
  return at::kComplexFloat;
}
 
template <> inline constexpr torch::Dtype dtype<std::complex<double>>() {
  return at::kComplexDouble;
}
 
inline int guess_device_index() {
#ifdef IGANET_WITH_MPI
  int rank;
  MPI_Comm_rank(MPI_COMM_WORLD, &rank);
  return rank %
         utils::getenv("IGANET_DEVICE_COUNT", (torch::cuda::is_available()
                                                  ? torch::cuda::device_count()
                                                  : (torch::xpu::is_available() ? torch::xpu::device_count() ? 1)));
#else
  return 0;
#endif
}
 
template <typename real_t>
class Options : private iganet::utils::FullQualifiedName {
public:
  Options()
      : options_(
            torch::TensorOptions()
                .dtype(::iganet::dtype<real_t>())
                .device_index(utils::getenv("IGANET_DEVICE_INDEX",
                                            iganet::guess_device_index()))
                .device(
                    (utils::getenv("IGANET_DEVICE", std::string{}) == "CPU")
                        ? torch::kCPU
                    : (utils::getenv("IGANET_DEVICE", std::string{}) == "CUDA")
                        ? torch::kCUDA
                    : (utils::getenv("IGANET_DEVICE", std::string{}) == "HIP")
                        ? torch::kHIP
                    : (utils::getenv("IGANET_DEVICE", std::string{}) == "MPS")
                        ? torch::kMPS
                    : (utils::getenv("IGANET_DEVICE", std::string{}) == "XLA")
                        ? torch::kXLA
                    : (utils::getenv("IGANET_DEVICE", std::string{}) == "XPU")
                        ? torch::kXPU
                        : (torch::cuda::is_available() ? torch::kCUDA
                                                       : (torch::xpu::is_available() ? torch::kXPU : torch::kCPU)))) {}
 
  explicit Options(torch::TensorOptions &&options)
      : options_(options.dtype(::iganet::dtype<real_t>())) {}
 
  operator torch::TensorOptions() const { return options_; }
 
  torch::Device device() const noexcept { return options_.device(); }
 
  int32_t device_index() const noexcept { return options_.device_index(); }
 
  torch::Dtype dtype() const noexcept { return ::iganet::dtype<real_t>(); }
 
  torch::Layout layout() const noexcept { return options_.layout(); }
 
  bool requires_grad() const noexcept { return options_.requires_grad(); }
 
  bool pinned_memory() const noexcept { return options_.pinned_memory(); }
 
  bool is_sparse() const noexcept { return options_.is_sparse(); }
 
  Options<real_t> device(torch::Device device) const noexcept {
    return Options(options_.device(device));
  }
 
  Options<real_t> device_index(int16_t device_index) const noexcept {
    return Options(options_.device_index(device_index));
  }
 
  template <typename other_t> Options<other_t> dtype() const noexcept {
    return Options<other_t>(options_.dtype(::iganet::dtype<other_t>()));
  }
 
  Options<real_t> layout(torch::Layout layout) const noexcept {
    return Options(options_.layout(layout));
  }
 
  Options<real_t> requires_grad(bool requires_grad) const noexcept {
    return Options(options_.requires_grad(requires_grad));
  }
 
  Options<real_t> pinned_memory(bool pinned_memory) const noexcept {
    return Options(options_.pinned_memory(pinned_memory));
  }
 
  Options<real_t>
  memory_format(torch::MemoryFormat memory_format) const noexcept {
    return Options(options_.memory_format(memory_format));
  }
 
  using value_type = real_t;
 
  inline virtual void
  pretty_print(std::ostream &os = Log(log::info)) const noexcept override {
    os << name() << "(\noptions = " << options_ << "\n)";
  }
 
private:
  const torch::TensorOptions options_;
};
 
template <typename real_t>
inline std::ostream &operator<<(std::ostream &os, const Options<real_t> &obj) {
  obj.pretty_print(os);
  return os;
}
 
template <typename real_t>
class Options<Options<real_t>> : public Options<real_t> {
  using Options<real_t>::Options;
};
 
} // namespace iganet