Bug fix(distance): correct CK pairwise distance on RDNA3 (gfx1100) WMMA

[zihaomu]

Summary

The Composable-Kernel (CK) pairwise-distance path was both unreachable and, once reached, numerically wrong on RDNA3 (gfx11) GPUs. This PR fixes two independent issues so that cuvsPairwiseDistance (L2Expanded / Cosine) correctly and efficiently uses RDNA3 WMMA.

Bugs

1. Dead-code 4GB guard (dispatch_ck.cuh)

The buffer-descriptor guard compared the theoretical maximum batch (mb_max * nb_max) against the 4GB limit instead of the effective per-batch output size min(mb_max, m) * min(nb_max, n). Since the theoretical product always exceeds 4GB, the guard always tripped and every problem size silently fell back to the SM60 kernel — the CK pairwise path was never executed.

2. Wrong warp-gemm for the C tile distribution (ck/epilogue.hpp)

The CShuffle epilogue selected the warp-gemm that defines the C tile-distribution using ODataType (the output type, float) rather than the GEMM input type.

• On CDNA, <float,float,float> and <half,half,float> resolve to the same MFMA C layout, so the bug is invisible.
• On RDNA3, <float,float,float> selects the MFMA warp-gemm (wave64, 4 C regs/thread) while the pipeline accumulator uses the fp16 WMMA warp-gemm (wave32, 8 interleaved C regs/thread). The mismatch made slice_acc_tile copy only 4 of 8 C registers, zeroing/scrambling half of the output rows.

Fix

• Thread the GEMM input type (DataT) into PairwiseDistanceCkEpilogueProblem and use it (instead of ODataType) to select the warp-gemm, so the epilogue C distribution matches the pipeline accumulator on both CDNA and RDNA.
• Guard on the effective per-batch dims so the CK path is actually reached.

3 files changed, +16 / −4.

Architecture-agnostic — no behavior change on CDNA. Selecting the warp-gemm by the input type is the correct choice on every architecture (the epilogue must match the pipeline accumulator). WarpGemmDispatcher already dispatches per-arch — WMMA on gfx11/gfx12, MFMA on gfx9 — so this simply lets that existing selection work rather than hard-coding anything. On CDNA the fp16 and fp32 MFMA 16×16 warp-gemms share the same C accumulator layout (which is why the previous ODataType code worked there), so the change is behavior-preserving on CDNA and only fixes the RDNA3 (WMMA ≠ MFMA) case.

Validation — Radeon PRO W7900 (gfx1100), ROCm 7.2.4

Repo test suite (DISTANCE_TEST): passes on gfx1100 — all L2Expanded + Cosine cases (101/101), including the fp16 half instances (DistanceEucExpTestH, DistanceExpCosH, and their XequalY variants) that now exercise the CK WMMA path. Differential check: reverting only the epilogue fix (keeping the guard so CK stays engaged) makes those fp16 cases fail — confirming the tests actually hit the CK path and that the fix is load-bearing. fp32/f64 continue on the SM60 path unchanged.

Independent harness: fp16 L2Expanded & Cosine also match a CPU fp64 reference exactly (ratio=1.0000, nbad=0) across sizes {16, 64, 100, 128, 256, 512}².

Performance (fp16 L2Expanded):

Size	SM60 (before)	CK WMMA (after)	Speedup
2048×2048×512	1.95 ms / 2204 GFLOP/s	0.182 ms / 23550 GFLOP/s	10.7×
4096×4096×512	8.99 ms / 1911 GFLOP/s	0.471 ms / 36501 GFLOP/s	19.1×

Scope

Only cuvsPairwiseDistance (L2Expanded / Cosine) is affected; brute-force / CAGRA / IVF do not use this dispatch.

[kevinjosephamd]

Thanks @zihaomu for this contribution. I might be missing something but it looks like the CK path is currently only enabled for GCN architectures that start with the gfx9 prefix. https://github.com/AMD-Ecosystem/hipVS/blob/release/rocmds-26.03/cpp/src/distance/detail/pairwise_matrix/dispatch_ck.cuh#L228-L232. Were there other local changes that forced control to take the CK path on your RDNA system? Should we relax that check?