Extract Metal shaders to standalone .metal file with kernel wrapper class (#5167)

Summary:
## Summary
Before submitting PRs for IVFFlat, IVFPQ, IVFSQ, and BinaryFlat implementations, I wanted to run this by you guys to see if this is a good approach for organizing the Metal shader code. Currently the MSL source is embedded as an inline string in the Objective C++ file, this extracts it into a standalone `.metal` file with a wrapper class, which will be the foundation for all the upcoming index type PRs. This also upgrades the top-k kernel from a single-thread insertion sort to a 256-thread parallel bitonic sort for better GPU utilization.

- Extract inline MSL shader source from `MetalFlatKernels.mm` into a standalone`MetalDistance.metal` file for proper syntax highlighting, easier editing, and reviewability
- Add `MetalKernels` wrapper class that loads a precompiled `.metallib`, caches pipeline states, and provides typed `encode*()` dispatch methods (previously the library was recompiled from source on every search call)
- Precompile `.metal` → `.metallib` via `xcrun metal`/`xcrun metallib` in CMake, catching shader errors at build time instead of runtime
- Upgrade distance kernels from naive per-element to tiled GEMM-style (32x32 tiles with shared memory) and upgrades top-k to a parallel threadgroup-based selection kernel (256-thread parallel bitonic sort)
- Foundation for future index types (IVFFlat, IVFPQ, etc.) to add kernels to the same `.metal` file and methods to `MetalKernels`

## Changes
- **New:** `MetalDistance.metal`, `l2_squared_matrix`, `ip_matrix`, `topk_threadgroup_{32,64,128,256}`
- **New:** `MetalKernels.h/.mm` runtime shader loading, pipeline caching, typed `encode*()` dispatch methods, per-device singleton
- **Modified:** `MetalFlatKernels.mm` replaced inline MSL + manual dispatch with `MetalKernels` calls
- **Modified:** `CMakeLists.txt` — added `MetalKernels.mm` to build, added custom commands to precompile `.metal` → `.air` → `.metallib`

## Build and test
```bash
cmake -B build \
  -DFAISS_ENABLE_METAL=ON \
  -DFAISS_ENABLE_GPU=OFF \
  -DFAISS_ENABLE_PYTHON=OFF \
  -DBUILD_TESTING=ON \
  -DCMAKE_BUILD_TYPE=Release \
  -DCMAKE_PREFIX_PATH="$(brew --prefix libomp)" \
  .
cmake --build build --target faiss_metal TestMetalIndexFlat -j$(sysctl -n hw.logicalcpu)
cd build && ctest -R TestMetalIndexFlat --output-on-failure
```

Pull Request resolved: #5167

Reviewed By: alibeklfc

Differential Revision: D103979808

Pulled By: mnorris11

fbshipit-source-id: 9583d6767e589ad74f7aa58141cd2e9034e5698b

Author: Evandabest

faiss/gpu_metal/CMakeLists.txt

@@ -11,6 +11,7 @@
 set(FAISS_METAL_SRC
   MetalResources.mm
   MetalIndex.mm
+  MetalKernels.mm
   MetalFlatKernels.mm
   MetalIndexFlat.mm
   StandardMetalResources.mm
@@ -44,3 +45,30 @@ set_target_properties(faiss_metal PROPERTIES
   OBJCXX_STANDARD 17
   OBJCXX_STANDARD_REQUIRED ON
 )
+
+# Pre-compile Metal shaders: .metal -> .air -> .metallib
+find_program(XCRUN xcrun REQUIRED)
+set(METAL_SOURCE ${CMAKE_CURRENT_SOURCE_DIR}/MetalDistance.metal)
+set(METAL_AIR ${CMAKE_CURRENT_BINARY_DIR}/MetalDistance.air)
+set(METAL_LIB ${CMAKE_CURRENT_BINARY_DIR}/MetalDistance.metallib)
+
+add_custom_command(
+  OUTPUT ${METAL_AIR}
+  COMMAND ${XCRUN} -sdk macosx metal -c ${METAL_SOURCE} -o ${METAL_AIR}
+  DEPENDS ${METAL_SOURCE}
+  COMMENT "Compiling MetalDistance.metal -> .air"
+)
+
+add_custom_command(
+  OUTPUT ${METAL_LIB}
+  COMMAND ${XCRUN} -sdk macosx metallib ${METAL_AIR} -o ${METAL_LIB}
+  DEPENDS ${METAL_AIR}
+  COMMENT "Linking MetalDistance.air -> .metallib"
+)
+
+add_custom_target(metal_shaders DEPENDS ${METAL_LIB})
+add_dependencies(faiss_metal metal_shaders)
+
+target_compile_definitions(faiss_metal PRIVATE
+  FAISS_METALLIB_PATH="${METAL_LIB}"
+)

faiss/gpu_metal/MetalDistance.metal

@@ -0,0 +1,241 @@
+// @lint-ignore-every LICENSELINT
+/**
+ * Copyright (c) Meta Platforms, Inc. and its affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ *
+ * Metal Shading Language kernels for distance computation and top-k selection.
+ *
+ * Kernel organization:
+ * - Distance kernels: l2_squared_matrix, ip_matrix (tiled GEMM-style)
+ * - Top-k selection: topk_threadgroup_K (parallel bitonic sort, K <= 256)
+ */
+
+#include <metal_stdlib>
+using namespace metal;
+
+kernel void l2_squared_matrix(
+    device const float* queries [[buffer(0)]],
+    device const float* vectors [[buffer(1)]],
+    device float* distances [[buffer(2)]],
+    device const uint* params [[buffer(3)]],
+    uint2 tgid [[threadgroup_position_in_grid]],
+    uint2 ltid [[thread_position_in_threadgroup]]
+) {
+    constexpr uint TILE_M = 32;
+    constexpr uint TILE_N = 32;
+    constexpr uint TILE_K = 16;
+    constexpr uint TG_THREADS = 256;
+
+    uint nq = params[0], nb = params[1], d = params[2];
+    uint row0 = tgid.y * TILE_M;
+    uint col0 = tgid.x * TILE_N;
+    uint ty = ltid.y, tx = ltid.x;
+    uint tid = ty * 16 + tx;
+
+    float acc00 = 0.0f, acc01 = 0.0f, acc10 = 0.0f, acc11 = 0.0f;
+
+    threadgroup float tgQ[TILE_M * TILE_K];
+    threadgroup float tgV[TILE_N * TILE_K];
+
+    for (uint dk = 0; dk < d; dk += TILE_K) {
+        uint kLen = min(TILE_K, d - dk);
+
+        for (uint i = tid; i < TILE_M * TILE_K; i += TG_THREADS) {
+            uint mr = i / TILE_K, mk = i % TILE_K;
+            uint gRow = row0 + mr;
+            tgQ[i] = (gRow < nq && mk < kLen) ? queries[gRow * d + dk + mk] : 0.0f;
+        }
+        for (uint i = tid; i < TILE_N * TILE_K; i += TG_THREADS) {
+            uint mr = i / TILE_K, mk = i % TILE_K;
+            uint gCol = col0 + mr;
+            tgV[i] = (gCol < nb && mk < kLen) ? vectors[gCol * d + dk + mk] : 0.0f;
+        }
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+
+        for (uint kk = 0; kk < TILE_K; kk++) {
+            float q0 = tgQ[(ty * 2) * TILE_K + kk];
+            float q1 = tgQ[(ty * 2 + 1) * TILE_K + kk];
+            float v0 = tgV[(tx * 2) * TILE_K + kk];
+            float v1 = tgV[(tx * 2 + 1) * TILE_K + kk];
+            float d00 = q0 - v0; acc00 += d00 * d00;
+            float d01 = q0 - v1; acc01 += d01 * d01;
+            float d10 = q1 - v0; acc10 += d10 * d10;
+            float d11 = q1 - v1; acc11 += d11 * d11;
+        }
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+    }
+
+    uint r0 = row0 + ty * 2, r1 = r0 + 1;
+    uint c0 = col0 + tx * 2, c1 = c0 + 1;
+    if (r0 < nq && c0 < nb) distances[r0 * nb + c0] = acc00;
+    if (r0 < nq && c1 < nb) distances[r0 * nb + c1] = acc01;
+    if (r1 < nq && c0 < nb) distances[r1 * nb + c0] = acc10;
+    if (r1 < nq && c1 < nb) distances[r1 * nb + c1] = acc11;
+}
+
+kernel void ip_matrix(
+    device const float* queries [[buffer(0)]],
+    device const float* vectors [[buffer(1)]],
+    device float* distances [[buffer(2)]],
+    device const uint* params [[buffer(3)]],
+    uint2 tgid [[threadgroup_position_in_grid]],
+    uint2 ltid [[thread_position_in_threadgroup]]
+) {
+    constexpr uint TILE_M = 32;
+    constexpr uint TILE_N = 32;
+    constexpr uint TILE_K = 16;
+    constexpr uint TG_THREADS = 256;
+
+    uint nq = params[0], nb = params[1], d = params[2];
+    uint row0 = tgid.y * TILE_M;
+    uint col0 = tgid.x * TILE_N;
+    uint ty = ltid.y, tx = ltid.x;
+    uint tid = ty * 16 + tx;
+
+    float acc00 = 0.0f, acc01 = 0.0f, acc10 = 0.0f, acc11 = 0.0f;
+
+    threadgroup float tgQ[TILE_M * TILE_K];
+    threadgroup float tgV[TILE_N * TILE_K];
+
+    for (uint dk = 0; dk < d; dk += TILE_K) {
+        uint kLen = min(TILE_K, d - dk);
+
+        for (uint i = tid; i < TILE_M * TILE_K; i += TG_THREADS) {
+            uint mr = i / TILE_K, mk = i % TILE_K;
+            uint gRow = row0 + mr;
+            tgQ[i] = (gRow < nq && mk < kLen) ? queries[gRow * d + dk + mk] : 0.0f;
+        }
+        for (uint i = tid; i < TILE_N * TILE_K; i += TG_THREADS) {
+            uint mr = i / TILE_K, mk = i % TILE_K;
+            uint gCol = col0 + mr;
+            tgV[i] = (gCol < nb && mk < kLen) ? vectors[gCol * d + dk + mk] : 0.0f;
+        }
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+
+        for (uint kk = 0; kk < TILE_K; kk++) {
+            float q0 = tgQ[(ty * 2) * TILE_K + kk];
+            float q1 = tgQ[(ty * 2 + 1) * TILE_K + kk];
+            float v0 = tgV[(tx * 2) * TILE_K + kk];
+            float v1 = tgV[(tx * 2 + 1) * TILE_K + kk];
+            acc00 += q0 * v0; acc01 += q0 * v1;
+            acc10 += q1 * v0; acc11 += q1 * v1;
+        }
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+    }
+
+    uint r0 = row0 + ty * 2, r1 = r0 + 1;
+    uint c0 = col0 + tx * 2, c1 = c0 + 1;
+    if (r0 < nq && c0 < nb) distances[r0 * nb + c0] = acc00;
+    if (r0 < nq && c1 < nb) distances[r0 * nb + c1] = acc01;
+    if (r1 < nq && c0 < nb) distances[r1 * nb + c0] = acc10;
+    if (r1 < nq && c1 < nb) distances[r1 * nb + c1] = acc11;
+}
+
+// ============================================================
+//  Parallel threadgroup-based top-k (bitonic sort)
+//  One threadgroup (256 threads) per query, 4 candidates per thread = 1024.
+// ============================================================
+
+#define TOPK_THREADGROUP_VARIANT(K) \
+kernel void topk_threadgroup_##K( \
+    device const float* distances [[buffer(0)]], \
+    device float* outDistances [[buffer(1)]], \
+    device int* outIndices [[buffer(2)]], \
+    device const uint* params [[buffer(3)]], \
+    uint qi [[threadgroup_position_in_grid]], \
+    uint tid [[thread_position_in_threadgroup]] \
+) { \
+    constexpr uint TG_SIZE = 256; \
+    constexpr uint R = 4; \
+    constexpr uint CANDIDATES = TG_SIZE * R; \
+    threadgroup float tgDist[CANDIDATES]; \
+    threadgroup int tgIdx[CANDIDATES]; \
+    uint nq = params[0], nb = params[1], k = params[2], want_min = params[3]; \
+    if (qi >= nq || k == 0) return; \
+    const device float* row = distances + qi * nb; \
+    uint kk = min(k, nb); \
+    uint K_out = min((uint)K, kk); \
+    \
+    float localDist[R]; \
+    int localIdx[R]; \
+    uint localCount = 0; \
+    \
+    for (uint j = tid; j < nb; j += TG_SIZE) { \
+        float v = row[j]; \
+        if (localCount < R) { \
+            uint pos = localCount; \
+            while (pos > 0 && ((want_min && v < localDist[pos-1]) || (!want_min && v > localDist[pos-1]))) { \
+                localDist[pos] = localDist[pos-1]; \
+                localIdx[pos] = localIdx[pos-1]; \
+                pos--; \
+            } \
+            localDist[pos] = v; \
+            localIdx[pos] = (int)j; \
+            localCount++; \
+        } else { \
+            bool better = want_min ? (v < localDist[R-1]) : (v > localDist[R-1]); \
+            if (better) { \
+                uint pos = R - 1; \
+                while (pos > 0 && ((want_min && v < localDist[pos-1]) || (!want_min && v > localDist[pos-1]))) { \
+                    localDist[pos] = localDist[pos-1]; \
+                    localIdx[pos] = localIdx[pos-1]; \
+                    pos--; \
+                } \
+                localDist[pos] = v; \
+                localIdx[pos] = (int)j; \
+            } \
+        } \
+    } \
+    \
+    for (uint i = 0; i < R; i++) { \
+        uint idx = tid * R + i; \
+        if (i < localCount) { \
+            tgDist[idx] = localDist[i]; \
+            tgIdx[idx] = localIdx[i]; \
+        } else { \
+            tgDist[idx] = want_min ? 1e38f : -1e38f; \
+            tgIdx[idx] = -1; \
+        } \
+    } \
+    threadgroup_barrier(mem_flags::mem_threadgroup); \
+    \
+    for (uint k2 = 2; k2 <= CANDIDATES; k2 *= 2) { \
+        for (uint j = k2 >> 1; j > 0; j >>= 1) { \
+            for (uint idx = tid; idx < CANDIDATES; idx += TG_SIZE) { \
+                uint partner = idx ^ j; \
+                if (partner < CANDIDATES && partner > idx) { \
+                    bool ascending = ((idx & k2) == 0); \
+                    bool partnerBetter = want_min \
+                        ? (tgDist[partner] < tgDist[idx] || (tgDist[partner] == tgDist[idx] && tgIdx[partner] < tgIdx[idx])) \
+                        : (tgDist[partner] > tgDist[idx] || (tgDist[partner] == tgDist[idx] && tgIdx[partner] < tgIdx[idx])); \
+                    bool idxBetter = want_min \
+                        ? (tgDist[idx] < tgDist[partner] || (tgDist[idx] == tgDist[partner] && tgIdx[idx] < tgIdx[partner])) \
+                        : (tgDist[idx] > tgDist[partner] || (tgDist[idx] == tgDist[partner] && tgIdx[idx] < tgIdx[partner])); \
+                    bool swap = ascending ? partnerBetter : idxBetter; \
+                    if (swap) { \
+                        float td = tgDist[idx]; tgDist[idx] = tgDist[partner]; tgDist[partner] = td; \
+                        int ti = tgIdx[idx]; tgIdx[idx] = tgIdx[partner]; tgIdx[partner] = ti; \
+                    } \
+                } \
+            } \
+            threadgroup_barrier(mem_flags::mem_threadgroup); \
+        } \
+    } \
+    \
+    for (uint i = tid; i < K_out; i += TG_SIZE) { \
+        outDistances[qi * k + i] = tgDist[i]; \
+        outIndices[qi * k + i] = tgIdx[i]; \
+    } \
+    for (uint i = tid; i < k - K_out; i += TG_SIZE) { \
+        outDistances[qi * k + K_out + i] = want_min ? 1e38f : -1e38f; \
+        outIndices[qi * k + K_out + i] = -1; \
+    } \
+}
+
+TOPK_THREADGROUP_VARIANT(32)
+TOPK_THREADGROUP_VARIANT(64)
+TOPK_THREADGROUP_VARIANT(128)
+TOPK_THREADGROUP_VARIANT(256)
+#undef TOPK_THREADGROUP_VARIANT