ops/gemv.hpp

Go to the documentation of this file.

#pragma once
 
#include <uni20/linalg/backends/cpu/gemv.hpp>
#include <uni20/linalg/dispatch.hpp>
#include <uni20/linalg/operation_tags.hpp>
#include <uni20/tensor/concepts.hpp>
 
#if UNI20_BACKEND_BLAS
#include <uni20/linalg/backends/blas/gemv.hpp>
#endif
 
#include <utility>
 
namespace uni20::linalg
{
 
template <class BackendSelector, uni20::MutableRankedTensorView<1> OutputTensor, class Scalar,
          uni20::RankedTensorView<2> MatrixTensor, uni20::RankedTensorView<1> InputTensor>
void gemv(BackendSelector&& selector, OutputTensor&& output, Scalar alpha, MatrixTensor const& matrix,
          InputTensor const& input, Scalar beta)
{
  auto output_span = uni20::mdspec_of(output);
  auto matrix_span = uni20::mdspec_of(matrix);
  auto input_span = uni20::mdspec_of(input);
  dispatch_kernel(std::forward<BackendSelector>(selector), gemv_op{}, output_span, alpha, matrix_span, input_span,
                  beta);
}
 
template <uni20::MutableRankedTensorView<1> OutputTensor, class Scalar, uni20::RankedTensorView<2> MatrixTensor,
          uni20::RankedTensorView<1> InputTensor>
void gemv(OutputTensor&& output, Scalar alpha, MatrixTensor const& matrix, InputTensor const& input, Scalar beta)
{
  auto selector = select_backend(gemv_op{}, output, matrix, input);
  gemv(selector, std::forward<OutputTensor>(output), alpha, matrix, input, beta);
}
 
} // namespace uni20::linalg