Optimize euclidian distance in raft refine phase

[anstellaire]

Initial issue

Originally written code (below) generated serial assembly and used strictly-ordered fadda instructions (at least on gcc & clang). That resulted in suboptimal performance.

for (size_t k = 0; k < dim; k++) {
  distance += DC::template eval<DistanceT>(query[k], row[k]);
}

Proposed solution

This PR provides optimized generic reduction using used partial vector sum (below), that helps vectorization but looses strcictly-ordered limitation.

template <typename DC, typename DistanceT, typename DataT>
DistanceT euclidean_distance_squared_generic(DataT const* a, DataT const* b, size_t n) {
  // vector register capacity in elements
  size_t constexpr vreg_len = (128 / 8) / sizeof(DistanceT);
  // unroll factor = vector register capacity * number of ports;
  size_t constexpr unroll_factor = vreg_len * 4;

  // unroll factor is a power of two
  size_t n_rounded = n & (0xFFFFFFFF ^ (unroll_factor - 1));
  DistanceT distance[unroll_factor] = {0};

  for (size_t i = 0; i < n_rounded; i += unroll_factor) {
    for (size_t j = 0; j < unroll_factor; ++j) {
      distance[j] += DC::template eval<DistanceT>(a[i + j], b[i + j]);
    }
  }

  for (size_t i = n_rounded; i < n; ++i) {
    distance[i] += DC::template eval<DistanceT>(a[i], b[i]);
  }

  for (size_t i = 1; i < unroll_factor; ++i) {
    distance[0] += distance[i];
  }

  return distance[0];
}

Also it adds NEON implementation with intrinsics which provided further speedup on certain test cases (can be removed if arch-specific code is undesired).

Results

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[cjnolet]

@anstellaire thank you for opening up a PR to improve the refinment APIs, but these have all been deprecated in RAFT and any further updates should be getting made in cuVS. Are you able to move your changes to cuVS library instead?

[anstellaire]

@cjnolet sure, will do.

[anstellaire]

The PR was opened in cuVS repository instead: rapidsai/cuvs#689.