address_space.hpp

Go to the documentation of this file.

// Copyright (c) 2022-2024 Simons Foundation
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0.txt
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
// Authors: Thomas Hahn, Miguel Morales, Nils Wentzell
 
#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;
 
} // 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 <MemoryArray A1, MemoryArray... 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 <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