Task: euclidean_distance_matrix
Category: Spatial--temporal algorithms
Difficulty: 1
Tags: pairwise_distance,embarrassingly_parallel,tiled_gemm_like
Source: kNN-CUDA: https://github.com/vincentfpgarcia/kNN-CUDA
Description: All-pairs squared Euclidean distance matrix between reference and query point sets.

Prompt template (YAML):
# Pairwise Squared Euclidean Distance Matrix -- Prompt Template

task_description: |
  Compute the pairwise squared Euclidean distance matrix between two sets of points.
  Given a reference set R of `ref_nb` points and a query set Q of `query_nb` points,
  each in `dim`-dimensional space, output the matrix
      dist[q, r] = sum_{k=0}^{dim-1} (Q[q,k] - R[r,k])^2
  for q in [0, query_nb), r in [0, ref_nb).

  This is the inner kernel of k-NN search, classification, clustering, and many
  spatial-distance pipelines. It is "embarrassingly parallel" but very memory-
  bound on GPU when implemented naively, so high-quality implementations use
  shared-memory tiling (similar to GEMM blocking).

  Based on kNN-CUDA (Garcia, Debreuve & Barlaud), `compute_distances` kernel.

interface: |
  ```c
  // Single end-to-end call. Do H2D + kernel + D2H and synchronize before returning.
  // Called multiple times (warmup + timed); must be idempotent.
  extern "C" void solution_compute(
      int ref_nb,
      int query_nb,
      int dim,
      const float* ref,
      const float* query,
      float* dist
  );
  ```

input_size_notes: |
  | Size   | ref_nb | query_nb | dim | output cells |
  |--------|--------|----------|-----|--------------|
  | small  | 128    | 128      | 32  | 16,384       |
  | medium | 512    | 512      | 128 | 262,144      |
  | large  | 2048   | 2048     | 256 | 4,194,304    |

  - All sizes are powers of two and multiples of typical tile sizes (16, 32, 64).
  - All inputs are float32 in row-major layout.

output_contract: |
  - `dist[r * query_nb + q]` = squared Euclidean distance from ref r to query q
  - Tolerance: atol=1e-2, rtol=1e-4
  - `solution_compute` must be idempotent and synchronize before returning

algorithm_background: |
  **Pairwise squared Euclidean distance**

  For each pair (q, r), compute
      dist[q, r] = sum_{k=0}^{dim-1} (Q[q,k] - R[r,k])^2

  Naive triple loop is O(query_nb * ref_nb * dim).

  The cells of `dist` are completely independent; the dot-product structure
  along the `dim` axis lends itself to GEMM-style shared-memory blocking
  exactly as in dense matrix multiply: load BLOCK x BLOCK tiles of points
  cooperatively, accumulate partial SSDs in registers, sweep along `dim`.

hints_l2: |
  - Treat the computation like SGEMM with the inner accumulator
    being `(a-b)^2` instead of `a*b`.
  - Tile both the (query, ref) output plane and the `dim` reduction axis.
  - 16x16 or 32x32 thread blocks work well; load tiles into shared memory
    cooperatively, then accumulate squared differences in registers.
  - Pre-upload `ref` and `query` to device memory once inside `solution_compute`.

hints_l1: |
  ### 1. Data Transfer
  - In `solution_compute` setup phase: `cudaMalloc` and `cudaMemcpy` ref, query (host -> device).
  - Allocate output `d_dist` once, keep state local to each call.
  - In `solution_compute`: launch kernel, then copy `d_dist` -> host `dist`.

  ### 2. Tiled Kernel Skeleton
  ```cuda
  #define BD 16  // tile dim
  __global__ void ed_kernel(int ref_nb, int query_nb, int dim,
                            const float* ref, const float* query,
                            float* dist) {
      __shared__ float sQ[BD][BD];
      __shared__ float sR[BD][BD];
      int tx = threadIdx.x, ty = threadIdx.y;
      int q = blockIdx.y * BD + ty;
      int r = blockIdx.x * BD + tx;
      float ssd = 0.f;
      for (int k0 = 0; k0 < dim; k0 += BD) {
          // Cooperative load
          sQ[ty][tx] = (q < query_nb && k0+tx < dim) ? query[q*dim + k0+tx] : 0.f;
          sR[ty][tx] = (r < ref_nb   && k0+tx < dim) ? ref  [r*dim + k0+tx] : 0.f;
          __syncthreads();
          if (q < query_nb && r < ref_nb) {
              #pragma unroll
              for (int k = 0; k < BD; k++) {
                  float t = sQ[ty][k] - sR[tx][k];
                  ssd += t * t;
              }
          }
          __syncthreads();
      }
      if (q < query_nb && r < ref_nb)
          dist[q*ref_nb + r] = ssd;
  }
  ```

  ### 3. Launch Configuration
  - Grid: `((ref_nb+BD-1)/BD, (query_nb+BD-1)/BD)`
  - Block: `(BD, BD)`


CPU reference implementation (C):
// cpu_reference.c — Pairwise squared Euclidean distance matrix (CPU baseline)
//
// Faithfully ported from kNN-CUDA (vincentfpgarcia/kNN-CUDA, code/knncuda.cu)
// — specifically the `compute_distances` global-memory kernel that computes
// the SSD (sum of squared differences) matrix between two point sets.
//
// Layout matches the original kNN-CUDA code:
//   * `ref`   stored column-major: ref[k * ref_pitch + r]
//             where k is the dim-row index and r is the point column index
//   * `query` stored column-major: query[k * query_pitch + q]
//   * `dist`  stored row-by-r:    dist[r * query_pitch + q]
//             i.e. one row per ref point, one column per query point.
//   * pitches equal nb (no padding) since input sizes are powers of 2.
//
// Reference: Garcia, Debreuve, Barlaud, "Fast k Nearest Neighbor Search using
// GPU", CVPRW 2008; Garcia et al., "k-NN search: fast GPU-based
// implementations", ICIP 2010.
//
// NO file I/O, NO main(). All I/O handled by task_io_cpu.c.

#include <stdlib.h>

// ===== Module-level state =====
static int          g_ref_nb;
static int          g_query_nb;
static int          g_dim;
static const float* g_ref;     // column-major [dim, ref_nb]
static const float* g_query;   // column-major [dim, query_nb]

// ===== Function ported from knncuda.cu::compute_distances =====
// Each "thread" (r, q) computes ssd over the dim axis, exactly like the
// CUDA kernel which assigns one thread per output cell.
// Output: dist[r * query_pitch + q] (matches kernel's write at line 89).
static void compute_distancesCpu(int          ref_nb,
                                 int          ref_pitch,
                                 int          query_nb,
                                 int          query_pitch,
                                 int          height,            // == dim
                                 const float* ref,
                                 const float* query,
                                 float*       dist)
{
    int r, q, k;
    for (r = 0; r < ref_nb; r++) {
        for (q = 0; q < query_nb; q++) {
            float ssd = 0.0f;
            for (k = 0; k < height; k++) {
                float a = ref  [k * ref_pitch   + r];
                float b = query[k * query_pitch + q];
                float tmp = a - b;
                ssd += tmp * tmp;
            }
            dist[r * query_pitch + q] = ssd;
        }
    }
}

// ===== Public interface =====

static void _orbench_old_init(int          ref_nb,
                   int          query_nb,
                   int          dim,
                   const float* ref,
                   const float* query)
{
    g_ref_nb   = ref_nb;
    g_query_nb = query_nb;
    g_dim      = dim;
    g_ref      = ref;
    g_query    = query;
}

static void _orbench_old_compute(int    ref_nb,
                      int    query_nb,
                      int    dim,
                      float* dist)
{
    // pitches == nb (no padding for power-of-two sizes)
    compute_distancesCpu(ref_nb, ref_nb, query_nb, query_nb, dim,
                         g_ref, g_query, dist);
}

// ── Unified compute_only wrapper (auto-migrated) ──
void solution_compute(int ref_nb, int query_nb, int dim, const float* ref, const float* query, float* dist) {
    _orbench_old_init(ref_nb, query_nb, dim, ref, query);
    _orbench_old_compute(ref_nb, query_nb, dim, dist);
}