tools.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 "../blas/tools.hpp"
#include "../exceptions.hpp"
#include "../mem/address_space.hpp"
#include "../traits.hpp"
 
#include <algorithm>
#include <cmath>
#include <concepts>
#include <cstdint>
#include <string_view>
#include <type_traits>
#include <utility>
 
namespace nda::tensor {
 
  // Import tools from the blas_lapack namespace.
  using namespace nda::blas_lapack;


#if defined(NDA_HAVE_CUDA) && defined(NDA_HAVE_CUTENSOR)
  static constexpr bool have_cutensor = true;
#else
  static constexpr bool have_cutensor = false;
#endif // NDA_HAVE_CUTENSOR

#ifdef NDA_HAVE_TBLIS
  static constexpr bool have_tblis = true;
#else
  static constexpr bool have_tblis = false;
#endif // NDA_HAVE_TBLIS

  template <BlasArrayOrConj A>
  using data_ptr_t = decltype(get_array(std::declval<A>()).data());

  enum class binary_op : std::uint8_t { SUM, PROD, SUM_ABS, MAX, MAX_ABS, MIN, MIN_ABS, NORM_2 };
 
  // clang-format off
  enum class unary_op : std::uint8_t {
    IDENTITY, SQRT, RELU, CONJ, RCP, SIGMOID, TANH, EXP, LOG, ABS, NEG,
    SIN, COS, TAN, SINH, COSH, ASIN, ACOS, ATAN, ASINH, ACOSH, ATANH,
    CEIL, FLOOR, MISH, SWISH, SOFT_PLUS, SOFT_SIGN
  };
  // clang-format on
 
  namespace detail {
 
    // Apply a nda::tensor::binary_op to two scalar operands.
    template <typename T>
    T apply_binary(binary_op op, T x, T y) {
      switch (op) {
        case binary_op::SUM: return x + y;
        case binary_op::PROD: return x * y;
        case binary_op::SUM_ABS: return std::abs(x) + std::abs(y);
        case binary_op::MAX_ABS: return std::max(std::abs(x), std::abs(y));
        case binary_op::MIN_ABS: return std::min(std::abs(x), std::abs(y));
        case binary_op::NORM_2: return std::sqrt(std::norm(x) + std::norm(y));
        case binary_op::MAX:
        case binary_op::MIN:
          if constexpr (!is_complex_v<T>) {
            return (op == binary_op::MAX ? std::max(x, y) : std::min(x, y));
          } else {
            NDA_RUNTIME_ERROR << "nda::tensor: binary_op::MAX/MIN are unsupported for complex value types";
          }
      }
      return T{}; // unreachable
    }
 
  } // namespace detail

  template <typename T>
  struct tensor_view {
    using value_type = T;

    T *data = nullptr;

    const long *extents = nullptr;

    const long *strides = nullptr;

    int ndim = 0;

    unary_op op = unary_op::IDENTITY;

    tensor_view() = default;

    tensor_view(T *p) : data(p) {}

    template <BlasArray A>
      requires std::convertible_to<data_ptr_t<A>, T *>
    tensor_view(A &&a, unary_op op) // NOLINT
       : data(get_array(a).data()),
         extents(get_array(a).indexmap().lengths().data()),
         strides(get_array(a).indexmap().strides().data()),
         ndim(get_rank<decltype(get_array(a))>),
         op(op) {}

    template <BlasArrayOrConj A>
      requires std::convertible_to<data_ptr_t<A>, T *>
    tensor_view(A &&a) : tensor_view(get_array(a), is_conj_array_expr<A> ? unary_op::CONJ : unary_op::IDENTITY) {} // NOLINT

    template <typename U>
      requires(!std::same_as<U, T> && std::convertible_to<U *, T *>)
    tensor_view(tensor_view<U> tv) // NOLINT
       : data(tv.data), extents(tv.extents), strides(tv.strides), ndim(tv.ndim), op(tv.op) {}
  };
 
  // Deduction guide: tensor_view(array) deduces T from the data pointer of the (underlying) array.
  template <BlasArray A>
  tensor_view(A &&, unary_op) -> tensor_view<std::remove_pointer_t<data_ptr_t<A>>>;
 
  template <BlasArrayOrConj A>
  tensor_view(A &&) -> tensor_view<std::remove_pointer_t<data_ptr_t<A>>>;

  template <typename T>
  using const_tensor_view = tensor_view<const T>;

  inline void require_equal_indices(std::string_view idx_a, std::string_view idx_b, int rank, std::string_view op_name) {
    if (static_cast<int>(idx_a.size()) != rank || idx_a != idx_b) {
      NDA_RUNTIME_ERROR << "nda::tensor::" << op_name << ": fallback to nda operations requires identical index strings of length " << rank
                        << ": idx_a = '" << idx_a << "', idx_b = '" << idx_b << "'";
    }
  }

  template <int R>
    requires(R >= 0 && R <= 26)
  std::string_view default_index() {
    static const auto arr = []() constexpr {
      std::array<char, R> s{};
      for (int i = 0; i < R; ++i) s[i] = static_cast<char>('a' + i);
      return s;
    }();
 
    return {arr.data(), arr.size()};
  }

 
} // namespace nda::tensor