linalg/ops/gemm.hpp

Go to the documentation of this file.

#pragma once
 
#include <uni20/linalg/backends/cpu/gemm.hpp>
#include <uni20/linalg/dispatch.hpp>
#include <uni20/linalg/operation_scalar.hpp>
#include <uni20/linalg/operation_tags.hpp>
#include <uni20/linalg/precision.hpp>
#include <uni20/tensor/concepts.hpp>
#include <uni20/tensor/output.hpp>
#include <uni20/tensor/precision.hpp>
#if UNI20_HAS_MPLAPACK_FLOAT80
#include <uni20/linalg/backends/blas/binary80.hpp>
#endif
 
#if UNI20_BACKEND_BLAS || UNI20_HAS_MPLAPACK_MPREAL
#include <uni20/linalg/backends/blas/gemm.hpp>
#endif
 
#if UNI20_BACKEND_CUBLAS
#include <uni20/linalg/backends/cublas/gemm.hpp>
#endif
 
#include <utility>
 
namespace uni20::linalg
{
template <class BackendSelector, uni20::MutableRankedTensorView<2> OutputTensor, class Alpha, class Beta,
          uni20::RankedTensorView<2> LhsTensor, uni20::RankedTensorView<2> RhsTensor>
  requires(OperationScalar<Alpha, tensor_element_t<OutputTensor>> &&
           OperationScalar<Beta, tensor_element_t<OutputTensor>>)
void gemm(BackendSelector&& selector, OutputTensor&& output, Alpha alpha, LhsTensor const& lhs, RhsTensor const& rhs,
          Beta beta, std::optional<Precision> requested = {})
{
  using scalar_type = tensor_element_t<OutputTensor>;
  auto output_span = uni20::mdspec_of(output);
  auto lhs_span = uni20::mdspec_of(lhs);
  auto rhs_span = uni20::mdspec_of(rhs);
  auto p = requested ? *requested : common_default_precision(lhs, rhs);
  scalar_type const beta_scalar(beta);
  if (!requested && beta_scalar != scalar_type{}) p = common_precision(p, precision_of(output));
  detail::dispatch_at_precision<scalar_type>(std::forward<BackendSelector>(selector), gemm_op{}, p, output_span,
                                             scalar_type(alpha), lhs_span, rhs_span, beta_scalar);
  uni20::detail::record_output_precision(output, p);
}
 
template <uni20::MutableRankedTensorView<2> OutputTensor, class Alpha, class Beta, uni20::RankedTensorView<2> LhsTensor,
          uni20::RankedTensorView<2> RhsTensor>
  requires(OperationScalar<Alpha, tensor_element_t<OutputTensor>> &&
           OperationScalar<Beta, tensor_element_t<OutputTensor>>)
void gemm(OutputTensor&& output, Alpha alpha, LhsTensor const& lhs, RhsTensor const& rhs, Beta beta,
          std::optional<Precision> p = {})
{
  auto selector = select_backend(gemm_op{}, output, lhs, rhs);
  gemm(selector, std::forward<OutputTensor>(output), alpha, lhs, rhs, beta, p);
}
 
} // namespace uni20::linalg