Add new popops::topk API and initial implementation based on bitonic …

…sort.

Summary:
Ref T29476.

 * Top and bottom k (through 'largest' argument)
 * Ascending/descending/unsorted return
 * any n/k/batch size
 * only float and half data types supported currently with unsigned as a secondary tensor to permute
 * half data supported through a cast to float
 * many optimisations left on the table

Reviewers: #poplar_team, #framework_ip_review_-_any_oss_or_third-party_code_use_has_been_approved, samk

Reviewed By: #poplar_team, #framework_ip_review_-_any_oss_or_third-party_code_use_has_been_approved, samk

Maniphest Tasks: T29476

Differential Revision: https://phabricator.sourcevertex.net/D40484

include/poplibs_support/Algorithm.hpp

@@ -26,6 +26,11 @@ template <typename T, typename U> constexpr static auto ceildiv(T x, U y) {
   return (x + y - 1) / y;
 }
 
+/// ceildiv with 0/0 = 0
+template <typename T> static inline T ceildiv0(T q, T d) {
+  return q + d != 0 ? (q + d - 1) / d : 0;
+}
+
 template <typename T, typename U> constexpr static auto floordiv(T x, U y) {
   static_assert(std::is_unsigned<T>::value && std::is_unsigned<U>::value,
                 "Only valid for unsigned integral types");
@@ -79,6 +84,18 @@ constexpr static std::pair<T, T> balancedPartition(T n, T d) {
   return std::make_pair(a, b);
 }
 
+/// Compute the nth triangular number.
+template <typename T> static inline T nthTriangular(T n) {
+  static_assert(std::is_unsigned_v<T>);
+  return (n * (n + 1)) / 2;
+}
+
+/// Check if a number is a power of 2.
+template <typename T> static inline bool isPowerOf2(T n) {
+  static_assert(std::is_unsigned_v<T>);
+  return ((n & (n - 1)) == 0);
+}
+
 } // end namespace poplibs_support
 
 #endif // poplibs_support_Algorithm_hpp

include/popops/SortOrder.hpp

@@ -0,0 +1,24 @@
+// Copyright (c) 2021 Graphcore Ltd. All rights reserved.
+
+#ifndef _popops_SortOrder_hpp_
+#define _popops_SortOrder_hpp_
+
+#include <iosfwd>
+
+namespace popops {
+
+/// Defines a required order for sorting operations.
+enum class SortOrder {
+  /// No ordering is required.
+  NONE,
+  /// Sort in ascending order.
+  ASCENDING,
+  /// Sort in descending order.
+  DESCENDING
+};
+
+std::ostream &operator<<(std::ostream &os, const SortOrder &o);
+
+} // end namespace popops
+
+#endif // _popops_SortOrder_hpp_

include/popops/TopK.hpp

@@ -0,0 +1,111 @@
+// 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 {
+
+/** Parameters for topK* APIs
+ */
+struct TopKParams {
+  /// The number of outputs from the top k operation.
+  /// This must be less or equal the number of elements in the innermost
+  /// dimension of the tensor used as input to the operation.
+  unsigned k;
+  /// If true, return the top k largest elements. Otherwise return the
+  /// top k smallest elements.
+  bool largest;
+  /// The required ordering of elements in the resulting tensor.
+  SortOrder sortOrder;
+
+  TopKParams(unsigned k, bool largest, SortOrder sortOrder) noexcept;
+};
+
+std::ostream &operator<<(std::ostream &os, const TopKParams &p);
+
+/** Create an return a new tensor laid out optimally to be used as
+ *  an input to a topK operation with the given parameters.
+ *
+ *  \param graph        The Poplar graph to add the tensor to.
+ *  \param type         The Poplar type of elements in the returned tensor.
+ *  \param shape        The shape of the returned tensor.
+ *  \param params       The parameters of the top k that the returned tensor
+ *                      will be used as input to.
+ *  \param debugContext Optional debug information.
+ *
+ *  \returns A newly created tensor with shape \p shape and full tile mapping.
+ */
+poplar::Tensor createTopKInput(poplar::Graph &graph, const poplar::Type &type,
+                               const std::vector<std::size_t> &shape,
+                               const TopKParams &params,
+                               const poplar::DebugContext &debugContext = {});
+
+/** Return the top k values in the innermost dimension of a tensor.
+ *
+ *  \param graph        The Poplar graph to add the operation to.
+ *  \param prog         The Poplar sequence to add the operation to.
+ *  \param t            The tensor in which to find the top-k values in
+ *                      the innermost dimension.
+ *  \param params       The parameters of the top k.
+ *  \param debugContext Optional debug information.
+ *
+ *  \returns A tensor with the top k values found in the innermost dimension
+ *           of \p t.
+ */
+poplar::Tensor topK(poplar::Graph &graph, poplar::program::Sequence &prog,
+                    const poplar::Tensor &t, const TopKParams &params,
+                    const poplar::DebugContext &debugContext = {});
+
+/** Return the top k values in the innermost dimension of a tensor along
+ *  with the permutation of another tensor with respect to the values.
+ *
+ *  \param graph        The Poplar graph to add the operation to.
+ *  \param prog         The Poplar sequence to add the operation to.
+ *  \param key          The tensor in which to find the top-k values in
+ *                      the innermost dimension.
+ *  \param value        A tensor with the same shape as \p key for which to
+ *                      get the permutation with respect to \p key.
+ *  \param params       The parameters of the top k.
+ *  \param debugContext Optional debug information.
+ *
+ *  \returns A pair of tensors. The first contains the top k values found
+ *           in the innermost dimension of \p key. The second contains the
+ *           permutation of the tensor \p value with respect to the tensor
+ *           \p key.
+ */
+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 = {});
+
+/** Return the top k values in the innermost dimension of a tensor along
+ *  with the indices of those values in the input tensor in the innermost
+ *  dimension.
+ *
+ *  \param graph        The Poplar graph to add the operation to.
+ *  \param prog         The Poplar sequence to add the operation to.
+ *  \param t            The tensor in which to find the top-k values in
+ *                      the innermost dimension.
+ *  \param params       The parameters of the top k.
+ *  \param debugContext Optional debug information.
+ *
+ *  \returns A pair of tensors. The first contains the top k values found
+ *           in the innermost dimension of \p t. The second contains the
+ *           indices of those values in the innermost dimension of \p t in
+ *           the original input.
+ */
+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_

lib/poplibs_test/Util.cpp

@@ -199,6 +199,7 @@ bool checkIsClose(FPType a, FPType b, double relativeTolerance) {
 
 template bool checkIsClose<bool>(bool, bool, double);
 template bool checkIsClose<int>(int, int, double);
+template bool checkIsClose<unsigned>(unsigned, unsigned, double);
 template bool checkIsClose<float>(float, float, double);
 template bool checkIsClose<double>(double, double, double);