address_space.hpp

Go to the documentation of this file.

// Copyright (c) 2022--present, The Simons Foundation
// This file is part of TRIQS/nda and is licensed under the Apache License, Version 2.0.
// SPDX-License-Identifier: Apache-2.0
// See LICENSE in the root of this distribution for details.

 
#pragma once
 
#include "../concepts.hpp"
#include "../device.hpp"
 
#include <algorithm>
 
namespace nda {

  // Forward declarations
  template <char OP, ArrayOrScalar L, ArrayOrScalar R>
  struct expr;
  template <typename F, Array... As>
  struct expr_call;
  template <char OP, Array A>
  struct expr_unary;
  template <int R, typename F>
  class array_adapter;
 
} // namespace nda
 
namespace nda::mem {


  enum class AddressSpace { None, Host, Device, Unified }; // Do not change order!

  using AddressSpace::Device;

  using AddressSpace::Host;

  using AddressSpace::None;

  using AddressSpace::Unified;

  template <typename T>
  static constexpr AddressSpace get_addr_space = mem::None;
 
  // Specialization of nda::mem::get_addr_space for const types.
  template <typename T>
  static constexpr AddressSpace get_addr_space<T const> = get_addr_space<T>;
 
  // Specialization of nda::mem::get_addr_space for reference types.
  template <typename T>
  static constexpr AddressSpace get_addr_space<T &> = get_addr_space<T>;

  template <AddressSpace A1, AddressSpace A2 = None, AddressSpace... As>
  constexpr AddressSpace combine = []() {
    static_assert(!(A1 == Host && A2 == Device) && !(A1 == Device && A2 == Host),
                  "Error in nda::mem::combine: Cannot combine Host and Device address spaces");
    if constexpr (sizeof...(As) > 0) { return combine<std::max(A1, A2), As...>; }
    return std::max(A1, A2);
  }();

  template <Array A1, Array... As>
  constexpr AddressSpace common_addr_space = combine<get_addr_space<A1>, get_addr_space<As>...>;

  template <MemoryArray A>
  static constexpr AddressSpace get_addr_space<A> = A::storage_t::address_space;

  template <Handle H>
  static constexpr AddressSpace get_addr_space<H> = H::address_space;

  template <char OP, ArrayOrScalar L, ArrayOrScalar R>
  static constexpr AddressSpace get_addr_space<expr<OP, L, R>> = combine<get_addr_space<L>, get_addr_space<R>>;

  template <typename F, Array... As>
  static constexpr AddressSpace get_addr_space<expr_call<F, As...>> = combine<get_addr_space<As>...>;

  template <char OP, Array A>
  static constexpr AddressSpace get_addr_space<expr_unary<OP, A>> = get_addr_space<A>;

  template <int R, typename F>
  static constexpr AddressSpace get_addr_space<array_adapter<R, F>> = Host;

  template <AddressSpace... AdrSpcs>
  static const auto check_adr_sp_valid = []() {
    static_assert(((AdrSpcs != None) & ...), "Error in nda::mem::check_adr_sp_valid: Cannot use None address space");
    static_assert(nda::have_device or ((AdrSpcs == Host) & ...),
                  "Error in nda::mem::check_adr_sp_valid: Device address space requires compiling with GPU support.");
  };

  template <typename... Ts>
    requires(sizeof...(Ts) > 0)
  static constexpr bool on_host = ((get_addr_space<Ts> == mem::Host) and ...);

  template <typename... Ts>
    requires(sizeof...(Ts) > 0)
  static constexpr bool on_device = ((get_addr_space<Ts> == mem::Device) and ...);

  template <typename... Ts>
    requires(sizeof...(Ts) > 0)
  static constexpr bool on_unified = ((get_addr_space<Ts> == mem::Unified) and ...);

  template <typename A0, typename... A>
  static constexpr bool have_same_addr_space = ((get_addr_space<A0> == get_addr_space<A>) and ... and true);

  template <typename... Ts>
  static constexpr bool have_host_compatible_addr_space = ((on_host<Ts> or on_unified<Ts>) and ...);

  template <typename... Ts>
  static constexpr bool have_device_compatible_addr_space = ((on_device<Ts> or on_unified<Ts>) and ...);

  template <typename... Ts>
  static constexpr bool have_compatible_addr_space = (have_host_compatible_addr_space<Ts...> or have_device_compatible_addr_space<Ts...>);
 
  // Test various combinations of address spaces.
  static_assert(combine<None, None> == None);
  static_assert(combine<Host, Host> == Host);
  static_assert(combine<None, Host> == Host);
  static_assert(combine<Host, None> == Host);
 
  static_assert(combine<Device, Device> == Device);
  static_assert(combine<None, Device> == Device);
  static_assert(combine<Device, None> == Device);
 
  static_assert(combine<Device, Unified> == Unified);
  static_assert(combine<Unified, Device> == Unified);
  static_assert(combine<Host, Unified> == Unified);
  static_assert(combine<Unified, Host> == Unified);

 
} // namespace nda::mem