Program Listing for File experiment_helper.h

↰ Return to documentation for file (src/sparsebase/experiment/experiment_helper.h)

#include <string>
#include <unordered_map>
#include <vector>

#include "sparsebase/bases/reorder_base.h"
#include "sparsebase/experiment/experiment_type.h"
#include "sparsebase/format/coo.h"
#include "sparsebase/format/csc.h"
#include "sparsebase/format/csr.h"
#include "sparsebase/format/format.h"
#include "sparsebase/format/format_order_one.h"
#include "sparsebase/format/format_order_two.h"

#ifndef SPARSEBASE_PROJECT_EXPERIMENT_HELPER_H
#define SPARSEBASE_PROJECT_EXPERIMENT_HELPER_H

namespace sparsebase::experiment {

template <template <typename, typename, typename> typename FormatType,
          template <typename, typename, typename> typename ReaderType,
          typename IDType, typename NNZType, typename ValueType>
std::unordered_map<std::string, format::Format *> LoadFormat(
    std::vector<std::string> &file_names) {
  auto reader = ReaderType<IDType, NNZType, ValueType>(file_names[0]);
  format::FormatOrderTwo<IDType, NNZType, ValueType> *coo = reader.ReadCOO();
  FormatType<IDType, NNZType, ValueType> *format =
      coo->template Convert<FormatType>();
  std::unordered_map<std::string, format::Format *> r;
  r["format"] = format;
  return r;
}

template <template <typename, typename, typename> typename ReaderType,
          typename IDType, typename NNZType, typename ValueType>
std::unordered_map<std::string, format::Format *> LoadCSR(
    std::vector<std::string> &file_names) {
  auto reader = ReaderType<IDType, NNZType, ValueType>(file_names[0]);
  format::CSR<IDType, NNZType, ValueType> *csr = reader.ReadCSR();
  std::unordered_map<std::string, format::Format *> r;
  r["format"] = csr;
  return r;
}

template <template <typename, typename, typename> typename ReaderType,
          typename IDType, typename NNZType, typename ValueType>
std::unordered_map<std::string, format::Format *> LoadCOO(
    std::vector<std::string> &file_names) {
  auto reader = ReaderType<IDType, NNZType, ValueType>(file_names[0]);
  format::COO<IDType, NNZType, ValueType> *csr = reader.ReadCOO();
  std::unordered_map<std::string, format::Format *> r;
  r["format"] = csr;
  return r;
}

template <template <typename, typename, typename> typename ReaderType,
          typename IDType, typename NNZType, typename ValueType>
std::unordered_map<std::string, format::Format *> LoadCSC(
    std::vector<std::string> &file_names) {
  auto reader = ReaderType<IDType, NNZType, ValueType>(file_names[0]);
  format::COO<IDType, NNZType, ValueType> *coo = reader.ReadCOO();
  format::CSC<IDType, NNZType, ValueType> *csc =
      coo->template Convert<format::CSC>();
  std::unordered_map<std::string, format::Format *> r;
  r["format"] = csc;
  return r;
}

template <template <typename, typename, typename> typename ReorderType,
          typename ContextType, typename IDType, typename NNZType,
          typename ValueType>
void ReorderCSR(std::unordered_map<std::string, format::Format *> &data,
                std::any fparams, std::any params) {
  ContextType context;
  auto p = std::any_cast<
      typename ReorderType<IDType, NNZType, ValueType>::ParamsType>(params);
  auto *perm = bases::ReorderBase::Reorder<ReorderType>(
      p, data["format"]->AsAbsolute<format::CSR<IDType, NNZType, ValueType>>(),
      {&context}, true);
  auto *A_reordered = bases::ReorderBase::Permute2D<format::CSR>(
      perm,
      data["format"]->AsAbsolute<format::CSR<IDType, NNZType, ValueType>>(),
      {&context}, true);
  auto *A_csc = A_reordered->template Convert<format::CSR>();
  data["processed_format"] = A_csc;
}

inline void Pass(std::unordered_map<std::string, format::Format *> &data,
                 std::any fparams, std::any params) {
  data["processed_format"] = data["format"];
}

template <template <typename, typename, typename> typename ReorderType,
          template <typename, typename, typename> typename FormatType,
          typename ContextType, typename IDType, typename NNZType,
          typename ValueType>
void Reorder(std::unordered_map<std::string, format::Format *> &data,
             std::any fparams, std::any params) {
  ContextType context;
  auto p = std::any_cast<
      typename ReorderType<IDType, NNZType, ValueType>::ParamsType>(params);
  auto *perm = bases::ReorderBase::Reorder<ReorderType>(
      p, data["format"]->AsAbsolute<FormatType<IDType, NNZType, ValueType>>(),
      {&context}, true);
  auto *A_reordered = bases::ReorderBase::Permute2D<FormatType>(
      perm,
      data["format"]->AsAbsolute<FormatType<IDType, NNZType, ValueType>>(),
      {&context}, true);
  auto *A_csc = A_reordered->template Convert<format::CSR>();
  data["processed_format"] = A_csc;
}
}  // namespace sparsebase::experiment

#endif  // SPARSEBASE_PROJECT_EXPERIMENT_HELPER_H