TopK.hpp

Go to the documentation of this file.

// Copyright (c) 2020 Graphcore Ltd. All rights reserved.
#ifndef _popops_TopK_hpp_
#define _popops_TopK_hpp_
 
#include <poplar/Graph.hpp>
#include <poplar/Program.hpp>
#include <poplar/Tensor.hpp>
 
#include <popops/SortOrder.hpp>
 
namespace popops {
 
struct TopKParams {
  unsigned k;
  bool largest;
  SortOrder sortOrder;
  bool stableSort;
 
  TopKParams(unsigned k, bool largest, SortOrder sortOrder,
             bool stableSort = false) noexcept;
};
 
std::ostream &operator<<(std::ostream &os, const TopKParams &p);
 
poplar::Tensor createTopKInput(poplar::Graph &graph, const poplar::Type &type,
                               const std::vector<std::size_t> &shape,
                               const TopKParams &params,
                               const poplar::DebugContext &debugContext = {});
 
poplar::Tensor topK(poplar::Graph &graph, poplar::program::Sequence &prog,
                    const poplar::Tensor &t, const TopKParams &params,
                    const poplar::DebugContext &debugContext = {});
 
std::pair<poplar::Tensor, poplar::Tensor>
topKKeyValue(poplar::Graph &graph, poplar::program::Sequence &prog,
             const poplar::Tensor &keys, const poplar::Tensor &values,
             const TopKParams &params,
             const poplar::DebugContext &debugContext = {});
 
std::pair<poplar::Tensor, poplar::Tensor>
topKWithPermutation(poplar::Graph &graph, poplar::program::Sequence &prog,
                    const poplar::Tensor &t, const TopKParams &params,
                    const poplar::DebugContext &debugContext = {});
 
} // end namespace popops
 
#endif // _popops_TopK_hpp_