pairwise_argmin.h

Go to the documentation of this file.

#pragma once
 
#include <algorithm>
#include <array>
#include <cstddef>
#include <cstdint>
#include <type_traits>
 
#include "clustering/always_assert.h"
#include "clustering/math/defaults.h"
#include "clustering/math/detail/pairwise_argmin_outer.h"
#include "clustering/math/pairwise.h"
#include "clustering/math/thread.h"
#include "clustering/ndarray.h"
 
namespace clustering::math {

inline constexpr std::size_t pairwiseArgminChunkRows = 256;

[[nodiscard]] inline std::array<std::size_t, 2>
chunkedMaterializedScratchShape(std::size_t n, std::size_t k) noexcept {
  const std::size_t rows = (n < pairwiseArgminChunkRows) ? n : pairwiseArgminChunkRows;
  return {rows == 0 ? std::size_t{1} : rows, k == 0 ? std::size_t{1} : k};
}
 
namespace detail {

enum class ArgminPath : std::uint8_t { Fused, Materialized };

template <class T, Layout LX, Layout LC>
void pairwiseArgminMaterializedWithScratch(const NDArray<T, 2, LX> &X, const NDArray<T, 2, LC> &C,
                                           NDArray<std::int32_t, 1> &labels,
                                           NDArray<T, 1> &outMinSq, NDArray<T, 2> &distsScratch,
                                           Pool pool) {
  const std::size_t n = X.dim(0);
  const std::size_t k = C.dim(0);
  const std::size_t d = X.dim(1);
  if (n == 0 || k == 0) {
    return;
  }
  const std::size_t chunkCap = pairwiseArgminChunkRows;
  CLUSTERING_ALWAYS_ASSERT(distsScratch.dim(1) >= k);
  CLUSTERING_ALWAYS_ASSERT(distsScratch.dim(0) >= (n < chunkCap ? n : chunkCap));
 
  auto runChunk = [&](std::size_t iBase, std::size_t chunkRows, const auto &xChunk) noexcept {
    // The scratch tile may be wider than k when the caller pre-sized with k padded up; the
    // dispatch view narrows it back to (chunkRows, k) so the callee sees exactly one chunk.
    NDArray<T, 2> distsView = NDArray<T, 2>::borrow(distsScratch.data(), {chunkRows, k});
    pairwiseSqEuclidean(xChunk, C, distsView, pool);
 
    auto scanRange = [&](std::size_t lo, std::size_t hi) noexcept {
      for (std::size_t i = lo; i < hi; ++i) {
        const T *row = distsView.data() + (i * k);
        T bestVal = row[0];
        std::int32_t bestIdx = 0;
        for (std::size_t j = 1; j < k; ++j) {
          const T v = row[j];
          if (v < bestVal) {
            bestVal = v;
            bestIdx = static_cast<std::int32_t>(j);
          }
        }
        outMinSq(iBase + i) = bestVal;
        labels(iBase + i) = bestIdx;
      }
    };
 
    if (pool.shouldParallelize(chunkRows, 4, 2) && pool.pool != nullptr) {
      pool.pool
          ->submit_blocks(std::size_t{0}, chunkRows,
                          [&](std::size_t lo, std::size_t hi) { scanRange(lo, hi); })
          .wait();
    } else {
      scanRange(0, chunkRows);
    }
  };
 
  for (std::size_t iBase = 0; iBase < n; iBase += chunkCap) {
    const std::size_t chunkRows = (iBase + chunkCap <= n) ? chunkCap : (n - iBase);
    if constexpr (LX == Layout::Contig) {
      auto xChunk = NDArray<T, 2, Layout::Contig>::borrow(X.data() + (iBase * d), {chunkRows, d});
      runChunk(iBase, chunkRows, xChunk);
    } else {
      auto xChunk = X.slice(0, iBase, iBase + chunkRows);
      runChunk(iBase, chunkRows, xChunk);
    }
  }
}

template <class T, Layout LX, Layout LC>
void pairwiseArgminMaterialized(const NDArray<T, 2, LX> &X, const NDArray<T, 2, LC> &C,
                                NDArray<std::int32_t, 1> &labels, NDArray<T, 1> &outMinSq,
                                Pool pool) {
  const std::size_t n = X.dim(0);
  const std::size_t k = C.dim(0);
  if (n == 0 || k == 0) {
    return;
  }
 
  const auto shape = chunkedMaterializedScratchShape(n, k);
  NDArray<T, 2> distsScratch({shape[0], shape[1]});
  pairwiseArgminMaterializedWithScratch(X, C, labels, outMinSq, distsScratch, pool);
}

template <class T, Layout LX, Layout LC>
bool canUseFusedArgmin(const NDArray<T, 2, LX> &X, const NDArray<T, 2, LC> &C,
                       const NDArray<T, 1> &cSqNorms) noexcept {
#ifdef CLUSTERING_USE_AVX2
  if constexpr (std::is_same_v<T, float> && LX == Layout::Contig && LC == Layout::Contig) {
    const std::size_t n = X.dim(0);
    const std::size_t k = C.dim(0);
    const std::size_t d = X.dim(1);
    if (n == 0 || k == 0 || d == 0) {
      return false;
    }
    if (d > defaults::pairwiseArgminMaxD) {
      return false;
    }
    if (!X.template isAligned<32>() || !C.template isAligned<32>()) {
      return false;
    }
    if (!cSqNorms.isContiguous()) {
      return false;
    }
    return true;
  } else {
    (void)X;
    (void)C;
    (void)cSqNorms;
    return false;
  }
#else
  (void)X;
  (void)C;
  (void)cSqNorms;
  return false;
#endif
}
 
} // namespace detail

template <class T, Layout LX = Layout::Contig, Layout LC = Layout::Contig>
void pairwiseArgminSqEuclidean(const NDArray<T, 2, LX> &X, const NDArray<T, 2, LC> &C,
                               const NDArray<T, 1> &cSqNorms, NDArray<std::int32_t, 1> &labels,
                               NDArray<T, 1> &outMinDistSq, Pool pool) {
  static_assert(std::is_same_v<T, float> || std::is_same_v<T, double>,
                "pairwiseArgminSqEuclidean<T> requires T to be float or double");
 
  CLUSTERING_ALWAYS_ASSERT(labels.isMutable());
  CLUSTERING_ALWAYS_ASSERT(outMinDistSq.isMutable());
  CLUSTERING_ALWAYS_ASSERT(X.dim(1) == C.dim(1));
  CLUSTERING_ALWAYS_ASSERT(labels.dim(0) == X.dim(0));
  CLUSTERING_ALWAYS_ASSERT(outMinDistSq.dim(0) == X.dim(0));
  CLUSTERING_ALWAYS_ASSERT(cSqNorms.dim(0) == C.dim(0));
 
  const std::size_t n = X.dim(0);
  const std::size_t k = C.dim(0);
  if (n == 0 || k == 0) {
    return;
  }
 
#ifdef CLUSTERING_USE_AVX2
  if constexpr (std::is_same_v<T, float> && LX == Layout::Contig && LC == Layout::Contig) {
    if (detail::canUseFusedArgmin(X, C, cSqNorms)) {
      detail::pairwiseArgminOuterAvx2F32(X, C, cSqNorms, labels, outMinDistSq, pool);
      return;
    }
  }
#endif
 
  detail::pairwiseArgminMaterialized(X, C, labels, outMinDistSq, pool);
}
 
namespace detail {

template <class T, Layout LX = Layout::Contig, Layout LC = Layout::Contig>
ArgminPath pairwiseArgminSqEuclideanWithDispatchInfo(const NDArray<T, 2, LX> &X,
                                                     const NDArray<T, 2, LC> &C,
                                                     const NDArray<T, 1> &cSqNorms,
                                                     NDArray<std::int32_t, 1> &labels,
                                                     NDArray<T, 1> &outMinDistSq, Pool pool) {
  static_assert(std::is_same_v<T, float> || std::is_same_v<T, double>,
                "pairwiseArgminSqEuclideanWithDispatchInfo<T> requires T to be float or double");
 
  CLUSTERING_ALWAYS_ASSERT(labels.isMutable());
  CLUSTERING_ALWAYS_ASSERT(outMinDistSq.isMutable());
  CLUSTERING_ALWAYS_ASSERT(X.dim(1) == C.dim(1));
  CLUSTERING_ALWAYS_ASSERT(labels.dim(0) == X.dim(0));
  CLUSTERING_ALWAYS_ASSERT(outMinDistSq.dim(0) == X.dim(0));
  CLUSTERING_ALWAYS_ASSERT(cSqNorms.dim(0) == C.dim(0));
 
  const std::size_t n = X.dim(0);
  const std::size_t k = C.dim(0);
  if (n == 0 || k == 0) {
    return ArgminPath::Materialized;
  }
 
#ifdef CLUSTERING_USE_AVX2
  if constexpr (std::is_same_v<T, float> && LX == Layout::Contig && LC == Layout::Contig) {
    if (canUseFusedArgmin(X, C, cSqNorms)) {
      pairwiseArgminOuterAvx2F32(X, C, cSqNorms, labels, outMinDistSq, pool);
      return ArgminPath::Fused;
    }
  }
#endif
 
  pairwiseArgminMaterialized(X, C, labels, outMinDistSq, pool);
  return ArgminPath::Materialized;
}
 
} // namespace detail
 
} // namespace clustering::math