add.hpp

Go to the documentation of this file.

// Copyright (c) 2024--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 "./interface/cutensor_interface.hpp"
#include "./interface/tblis_interface.hpp"
#include "./tools.hpp"
#include "../mem/address_space.hpp"
#include "../traits.hpp"
 
#include <string_view>
#include <utility>
 
namespace nda::tensor {


  template <BlasArrayOrConj A, BlasArrayFor<A> B>
  void add(get_value_t<A> alpha, A const &a, std::string_view idx_a, get_value_t<A> beta, B &&b, std::string_view idx_b) { // NOLINT
    // compile-time checks
    constexpr bool run_on_device = mem::have_device_compatible_addr_space<A, B>;
    static_assert(!run_on_device || have_cutensor, "nda::tensor::add: cuTENSOR support is required");
    static_assert(run_on_device || have_tblis || get_rank<A> == get_rank<B>, "nda::tensor::add: host fallback requires identical ranks");
 
    // dispatch to backends
    if constexpr (run_on_device) {
      device::elementwise_binary(alpha, a, idx_a, beta, b, idx_b, b);
    } else if constexpr (have_tblis) {
      tblis::add(alpha, a, idx_a, beta, b, idx_b);
    } else {
      require_equal_indices(idx_a, idx_b, get_rank<A>, "add");
      b = alpha * a + beta * b;
    }
  }

  template <BlasArrayOrConj A, BlasArrayOrConjFor<A> B, BlasArrayFor<A> C>
  void add(get_value_t<A> alpha, A const &a, std::string_view idx_a, get_value_t<A> beta, B const &b, std::string_view idx_b, C &&c, // NOLINT
           std::string_view idx_c) {
    // compile-time checks
    constexpr bool run_on_device = mem::have_device_compatible_addr_space<A, B, C>;
    static_assert(!run_on_device || have_cutensor, "nda::tensor::add: cuTENSOR support is required");
    static_assert(run_on_device || have_tblis || (get_rank<A> == get_rank<B> && get_rank<A> == get_rank<C>),
                  "nda::tensor::add: host fallback requires identical ranks");
 
    // dispatch to backends
    if constexpr (run_on_device) {
      device::elementwise_binary(alpha, a, idx_a, beta, b, idx_b, c, binary_op::SUM);
    } else if constexpr (have_tblis) {
      tblis::add(beta, b, idx_b, get_value_t<A>{0}, c, idx_c);
      tblis::add(alpha, a, idx_a, get_value_t<A>{1}, c, idx_c);
    } else {
      require_equal_indices(idx_a, idx_b, get_rank<A>, "add");
      require_equal_indices(idx_b, idx_c, get_rank<A>, "add");
      c = alpha * a + beta * b;
    }
  }

  template <BlasArrayOrConj A, BlasArrayFor<A> B>
  void add(A const &a, std::string_view idx_a, B &&b, std::string_view idx_b) { // NOLINT
    add(get_value_t<A>{1}, a, idx_a, get_value_t<A>{0}, std::forward<B>(b), idx_b);
  }

  template <BlasArrayOrConj A, BlasArrayOrConjFor<A> B, BlasArrayFor<A> C>
  void add(A const &a, std::string_view idx_a, B const &b, std::string_view idx_b, C &&c, std::string_view idx_c) { // NOLINT
    add(get_value_t<A>{1}, a, idx_a, get_value_t<A>{1}, b, idx_b, std::forward<C>(c), idx_c);
  }

  template <BlasArrayOrConj A, BlasArrayFor<A> B>
  void add(get_value_t<A> alpha, A const &a, get_value_t<A> beta, B &&b) { // NOLINT
    add(alpha, a, default_index<get_rank<A>>(), beta, std::forward<B>(b), default_index<get_rank<B>>());
  }

  template <BlasArrayOrConj A, BlasArrayFor<A> B>
  void add(A const &a, B &&b) { // NOLINT
    add(get_value_t<A>{1}, a, default_index<get_rank<A>>(), get_value_t<A>{0}, std::forward<B>(b), default_index<get_rank<B>>());
  }

 
} // namespace nda::tensor