BlockSparseMatMul.hpp

Go to the documentation of this file.

// Copyright (c) 2020 Graphcore Ltd. All rights reserved.
#ifndef POPSPARSE_BLOCK_SPARSE_MATMUL_H
#define POPSPARSE_BLOCK_SPARSE_MATMUL_H
 
#include "BlockSparse.hpp"
#include <array>
#include <poplar/Graph.hpp>
 
namespace popsparse {
namespace experimental {
 
class BSMatMulImpl;
 
class BSMatMulParams {
 
public:
  BSMatMulParams(const std::array<int, 3> &dim,
                 const std::array<int, 3> &blockSize,
                 const std::vector<unsigned char> &rhsSparsity,
                 bool rhsNeedTranspose, poplar::Type inDataType,
                 poplar::Type outDataType, poplar::Type partialDataType,
                 unsigned numGroupsIn = 1);
 
  BSMatMulParams(const std::array<int, 3> &dim,
                 const std::array<int, 3> &blockSize,
                 const std::vector<unsigned char> &resSparsity,
                 poplar::Type inDataType, poplar::Type outDataType,
                 poplar::Type partialDataType,
                 SubBlockMask subBlockMask = SubBlockMask::None,
                 unsigned numGroupsIn = 1);
 
  BSMatMulParams(BSMatMulParams &&other);
 
  ~BSMatMulParams();
 
  std::unique_ptr<BSMatMulImpl> impl;
};
 
poplar::Tensor createBSMatMulInputLHS(poplar::Graph &graph,
                                      const BSMatMulParams &bsMatMul,
                                      const poplar::DebugContext &debugContext,
                                      const poplar::OptionFlags &options = {});
 
poplar::Tensor createBSMatMulInputRHS(poplar::Graph &graph,
                                      const BSMatMulParams &bsMatMul,
                                      const poplar::DebugContext &debugContext,
                                      const poplar::OptionFlags &options = {});
 
/* [INTERNAL OPTIONS]
 *   * `memoryCycleRatio` Integer [=1]
 *
 *      This is used to compute the weight of a hypergraph node:
 *      w = memory_cycle_ratio * mem_weight +
 *          (1.0 - memory_cycle_ratio) * cycle_weight
 *      This may be only a temporary option.
 *
 *   * `partitionMethod` (block, block-naive, strip) [=block]
 *
 *      * **block:** The matrix multiply computation
 *        graph is created for each non-zero block and Zoltan
 *        is used to partition the graph.
 *
 *      * **block-naive:** The matrix multiply
 *        computation graph is created for each non-zero block
 *        and a greedy algorithm is used to partition the
 *        graph.
 *
 *      * **strip:** The graph is created for columns or rows.
 */
poplar::Tensor bsMatMul(poplar::Graph &graph,
                        const BSMatMulParams &bsMatMulParams,
                        poplar::program::Sequence &prog,
                        const poplar::Tensor &lhsMatrix,
                        const poplar::Tensor &rhsMatrix,
                        const poplar::OptionFlags &options = {},
                        const poplar::DebugContext &debugContext = {});
 
} // namespace experimental
} // namespace popsparse
 
#endif // POPSPARSE_BLOCK_SPARSE_MATMUL_H