BlockSparse.hpp

Go to the documentation of this file.

// Copyright (c) 2020 Graphcore Ltd. All rights reserved.
#ifndef POPSPARSE_BLOCK_SPARSE_H
#define POPSPARSE_BLOCK_SPARSE_H
 
#include <array>
#include <string>
#include <vector>
 
#include <poplar/Graph.hpp>
#include <poplar/Program.hpp>
#include <poplar/Tensor.hpp>
#include <poplar/Type.hpp>
 
namespace popsparse {
namespace experimental {
 
enum class SubBlockMask {
  None,
  ZeroUpperTriangle,
  ZeroLowerTriangle
};
 
poplar::Tensor bsSoftmax(poplar::Graph &graph, poplar::Tensor sparseTensor,
                         const std::array<int, 2> &dim,
                         const std::array<int, 2> &blockSize,
                         const std::vector<unsigned char> &sparsity,
                         SubBlockMask subBlockMaskType, unsigned numGroups,
                         poplar::program::Sequence &prog,
                         const poplar::DebugContext &debugContext = {});
 
void bsSoftmaxInPlace(poplar::Graph &graph, poplar::Tensor sparseTensor,
                      const std::array<int, 2> &dim,
                      const std::array<int, 2> &blockSize,
                      const std::vector<unsigned char> &sparsity,
                      SubBlockMask subBlockMaskType, unsigned numGroups,
                      poplar::program::Sequence &prog,
                      const poplar::DebugContext &debugContext = {});
 
poplar::Tensor bsSoftmaxGrad(poplar::Graph &graph, poplar::Tensor sparseOut,
                             poplar::Tensor sparseOutGrad,
                             const std::array<int, 2> &dim,
                             const std::array<int, 2> &blockSize,
                             const std::vector<unsigned char> &sparsity,
                             poplar::program::Sequence &prog,
                             const poplar::DebugContext &debugContext = {});
 
} // namespace experimental
} // namespace popsparse
 
#endif // POPSPARSE_BLOCK_SPARSE_H