sparse_mm_mkl_template Class — pytorch Architecture

Source Code

aten/src/ATen/native/mkl/SparseCsrLinearAlgebra.cpp lines 180–206

template <typename scalar_t>
static inline void sparse_mm_mkl_template(
    Tensor& res,
    const Tensor& col_indices,
    const Tensor& crow_indices,
    const Tensor& values,
    const Tensor& dense,
    const Tensor& t,
    const Scalar& alpha,
    const Scalar& beta,
    IntArrayRef size,
    IntArrayRef dense_size) {
  SparseCsrMKLInterface mkl_impl(
      col_indices.data_ptr<MKL_INT>(),
      crow_indices.data_ptr<MKL_INT>(),
      values.data_ptr<scalar_t>(),
      size[0],
      size[1]);
  mkl_impl.sparse_mm(
      res.data_ptr<scalar_t>(),
      dense.data_ptr<scalar_t>(),
      alpha.to<scalar_t>(),
      beta.to<scalar_t>(),
      size[0],
      size[1],
      dense_size[1]);
}