scale.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 "../exceptions.hpp"
#include "../mapped_functions.hpp"
#include "../mapped_functions.hxx"
#include "../mem/address_space.hpp"
#include "../traits.hpp"
 
namespace nda::tensor {


  template <BlasArray A>
  void scale(get_value_t<A> alpha, A &&a, unary_op op = unary_op::IDENTITY) { // NOLINT
    // compile-time checks
    constexpr bool run_on_device = mem::have_device_compatible_addr_space<A>;
    static_assert(!run_on_device || have_cutensor, "nda::tensor::scale: cuTENSOR support is required");
 
    // fold NEG into the scalar (alpha * (-x) == -alpha * x) so backends always see IDENTITY
    if (op == unary_op::NEG) {
      alpha = -alpha;
      op    = unary_op::IDENTITY;
    }
 
    // dispatch to backends
    if constexpr (run_on_device) {
      device::permute(alpha, tensor_view{a, op}, default_index<get_rank<A>>(), tensor_view{a, unary_op::IDENTITY}, default_index<get_rank<A>>());
    } else {
      // TBLIS only handles IDENTITY and CONJ; other ops fall through to the nda fallback below
      if constexpr (have_tblis) {
        if (op == unary_op::IDENTITY || op == unary_op::CONJ) {
          tblis::scale(alpha, tensor_view{a, op}, default_index<get_rank<A>>());
          return;
        }
      }
      switch (op) {
        case unary_op::IDENTITY: a = alpha * a; break;
        case unary_op::CONJ: a = alpha * nda::conj(a); break;
        case unary_op::ABS: a = alpha * nda::abs(a); break;
        case unary_op::SQRT: a = nda::map([alpha](auto x) { return alpha * std::sqrt(x); })(a); break;
        case unary_op::EXP: a = nda::map([alpha](auto x) { return alpha * std::exp(x); })(a); break;
        case unary_op::LOG: a = nda::map([alpha](auto x) { return alpha * std::log(x); })(a); break;
        case unary_op::RCP: a = nda::map([alpha](auto x) { return alpha / x; })(a); break;
        default:
          NDA_RUNTIME_ERROR << "nda::tensor::scale: unsupported unary_op on nda host fallback "
                               "(supported: IDENTITY, CONJ, NEG, SQRT, ABS, EXP, LOG, RCP)";
      }
    }
  }

 
} // namespace nda::tensor