SparsePartitioner.hpp

Go to the documentation of this file.

// Copyright (c) 2020 Graphcore Ltd. All rights reserved.
#ifndef _poplibs_popsparse_SparsePartitioner_hpp_
#define _poplibs_popsparse_SparsePartitioner_hpp_
 
#include <poplar/Interval.hpp>
#include <poplar/OptionFlags.hpp>
#include <poplar/Target.hpp>
#include <poplar/Type.hpp>
#include <popsparse/FullyConnected.hpp>
#include <popsparse/FullyConnectedParams.hpp>
#include <popsparse/MatMulParams.hpp>
#include <popsparse/SparseStorageFormats.hpp>
#include <string>
#include <vector>
 
namespace popsparse {
 
class PartitionerImpl;
 
namespace dynamic {
 
template <typename T> struct SparsityDataImpl {
  std::vector<std::size_t> metaInfo;
 
  std::vector<T> nzValues;
};
 
template <typename T> class Partitioner {
  std::string name;
 
public:
  const PartitionerImpl &getImpl() const { return *impl; }
 
  Partitioner(const FullyConnectedParams &params, const poplar::Type &dataType,
              const poplar::Target &target, const poplar::OptionFlags &options,
              PlanningCache *cache = {}, std::string name = "");
 
  Partitioner(const MatMulParams &params, const poplar::Type &dataType,
              const poplar::Target &target, const poplar::OptionFlags &options,
              PlanningCache *cache = {}, std::string name = "");
 
  ~Partitioner();
 
  SparsityDataImpl<T> createSparsityDataImpl(const CSCMatrix<T> &matrix_) const;
 
  SparsityDataImpl<T> createSparsityDataImpl(const CSRMatrix<T> &matrix_) const;
 
  SparsityDataImpl<T> createSparsityDataImpl(const COOMatrix<T> &matrix_) const;
 
  COOMatrix<T> sparsityDataImplToCOOMatrix(
      const SparsityDataImpl<T> &sparsityDataImpl) const;
 
  CSRMatrix<T> sparsityDataImplToCSRMatrix(
      const SparsityDataImpl<T> &sparsityDataImpl) const;
 
  CSCMatrix<T> sparsityDataImplToCSCMatrix(
      const SparsityDataImpl<T> &sparsityDataImpl) const;
 
  std::array<std::vector<std::size_t>, 3> getPlanPartitions(void) const;
 
  std::size_t getmetaInfoBucketElements(void) const;
 
private:
  std::unique_ptr<PartitionerImpl> impl;
};
 
} // namespace dynamic
} // namespace popsparse
#endif // _poplibs_popsparse_SparsePartitioner_hpp_