[BUG] l2_exp & KL distances got randomly round-off errors in half-float32 mixed precision on self-neighboring

[rhdong]

• This PR resolved this issue partially
• Considering the issue becomes more serious as the k larger, there's a need to discuss a perfect solution
  • It's hard to determine the thresholds empirically for get_clamp_precision. Please refer to the code: https://github.com/rapidsai/cuvs/blob/branch-25.02/cpp/src/distance/detail/distance_ops/l2_exp.cuh#L32

[rhdong]

• Reproduce code:

#include <iostream>
#include <cuda_fp16.h>
#include <cmath>
#include <cstdlib>
#include <ctime>
#include <iomanip>

__global__ void euclidean_distance_half(
    const __half* a, const __half* b, float* result, int n) {
    int idx = threadIdx.x + blockIdx.x * blockDim.x;

    __shared__ float sdata[256];

    float local_sum_a = 0.0f;
    float local_sum_b = 0.0f;
    float local_dot_product = 0.0f;

    for (int i = idx; i < n; i += blockDim.x * gridDim.x) {
        float af = __half2float(a[i]);
        float bf = __half2float(b[i]);
        local_sum_a += af * af;
        local_sum_b += bf * bf;
        local_dot_product += af * bf;
    }

    sdata[threadIdx.x] = local_sum_a;
    __syncthreads();

    for (int s = blockDim.x / 2; s > 0; s >>= 1) {
        if (threadIdx.x < s) {
            sdata[threadIdx.x] += sdata[threadIdx.x + s];
        }
        __syncthreads();
    }

    if (threadIdx.x == 0) {
        atomicAdd(result, sdata[0]);
    }

    __syncthreads();

    sdata[threadIdx.x] = local_sum_b;
    __syncthreads();

    for (int s = blockDim.x / 2; s > 0; s >>= 1) {
        if (threadIdx.x < s) {
            sdata[threadIdx.x] += sdata[threadIdx.x + s];
        }
        __syncthreads();
    }

    if (threadIdx.x == 0) {
        atomicAdd(result + 1, sdata[0]);
    }

    __syncthreads();

    sdata[threadIdx.x] = local_dot_product;
    __syncthreads();

    for (int s = blockDim.x / 2; s > 0; s >>= 1) {
        if (threadIdx.x < s) {
            sdata[threadIdx.x] += sdata[threadIdx.x + s];
        }
        __syncthreads();
    }

    if (threadIdx.x == 0) {
        atomicAdd(result + 2, sdata[0]);
    }
}

void test_euclidean_distance() {
    const int VECTOR_SIZE = 1024;
    const int TEST_COUNT = 1000;
    const float ERROR_THRESHOLD = 0.0001;

    __half* d_a;
    __half* d_b;
    float* d_result;
    float h_result[3];

    cudaMalloc(&d_a, VECTOR_SIZE * sizeof(__half));
    cudaMalloc(&d_b, VECTOR_SIZE * sizeof(__half));
    cudaMalloc(&d_result, 3 * sizeof(float));

    std::srand(std::time(nullptr));

    for (int test = 0; test < TEST_COUNT; ++test) {
        __half* h_a = new __half[VECTOR_SIZE];
        __half* h_b = new __half[VECTOR_SIZE];
        for (int i = 0; i < VECTOR_SIZE; ++i) {
            float random_value = static_cast<float>(std::rand()) / RAND_MAX;
            h_a[i] = __float2half(random_value);
            h_b[i] = h_a[i];
        }

        cudaMemcpy(d_a, h_a, VECTOR_SIZE * sizeof(__half), cudaMemcpyHostToDevice);
        cudaMemcpy(d_b, h_b, VECTOR_SIZE * sizeof(__half), cudaMemcpyHostToDevice);

        cudaMemset(d_result, 0, 3 * sizeof(float));

        int threads = 256;
        int blocks = (VECTOR_SIZE + threads - 1) / threads;
        euclidean_distance_half<<<blocks, threads>>>(d_a, d_b, d_result, VECTOR_SIZE);

        cudaMemcpy(h_result, d_result, 3 * sizeof(float), cudaMemcpyDeviceToHost);

        float distance = std::sqrt(h_result[0] + h_result[1] - 2 * h_result[2]);

        float error = std::abs(distance - 0.0f);
        if (error > ERROR_THRESHOLD) {
            std::cout << std::fixed << std::setprecision(7);
            std::cout << "Test " << test << " failed with error: " << error << std::endl;
            std::cout << "a^2: " << h_result[0] << ", b^2: " << h_result[1]
                      << ", dot: " << h_result[2] << ", distance: " << distance << std::endl;

            std::cout << "a: ";
            for (int i = 0; i < VECTOR_SIZE; ++i) {
                std::cout << __half2float(h_a[i]) << " ";
            }
            std::cout << "\nb: ";
            for (int i = 0; i < VECTOR_SIZE; ++i) {
                std::cout << __half2float(h_b[i]) << " ";
            }
            std::cout << std::endl;
        }

        delete[] h_a;
        delete[] h_b;
    }

    cudaFree(d_a);
    cudaFree(d_b);
    cudaFree(d_result);
}

int main() {
    test_euclidean_distance();
    return 0;
}