facebookresearch · mnorris11 · Jan 27, 2025
diff --git a/.github/actions/build_cmake/action.yml b/.github/actions/build_cmake/action.yml
@@ -23,12 +23,19 @@ runs:
       uses: conda-incubator/setup-miniconda@v3
       with:
         python-version: '3.11'
-        miniconda-version: latest
+        miniforge-version: latest # ensures conda-forge channel is used.
+        channels: conda-forge
+        conda-remove-defaults: 'true'
+        # Set to aarch64 if we're on arm64 because there's no miniforge ARM64 package, just aarch64.
+        # They are the same thing, just named differently.
+        architecture: ${{ runner.arch  == 'ARM64' && 'aarch64' || runner.arch }}
     - name: Configure build environment
       shell: bash
       run: |
         # initialize Conda
         conda config --set solver libmamba
+        # Ensure starting packages are from conda-forge.
+        conda list --show-channel-urls
         conda update -y -q conda
         echo "$CONDA/bin" >> $GITHUB_PATH
 
@@ -43,7 +50,7 @@ runs:
         if [ "${{ runner.arch }}" = "X64" ]; then
           # TODO: merge this with ARM64
           conda install -y -q -c conda-forge gxx_linux-64=14.2 sysroot_linux-64=2.17
-          conda install -y -q mkl=2023 mkl-devel=2023
+          conda install -y -q mkl=2022.2.1 mkl-devel=2022.2.1
         fi
 
         # no CUDA needed for ROCm so skip this
@@ -56,6 +63,7 @@ runs:
         elif [ "${{ inputs.cuvs }}" = "ON" ]; then
           conda install -y -q libcuvs=24.12 'cuda-version>=12.0,<=12.5' cuda-toolkit=12.4.1 gxx_linux-64=12.4 -c rapidsai -c conda-forge
         fi
+
         # install test packages
         if [ "${{ inputs.rocm }}" = "ON" ]; then
           : # skip torch install via conda, we need to install via pip to get
@@ -174,3 +182,8 @@ runs:
       with:
         name: test-results-arch=${{ runner.arch }}-opt=${{ inputs.opt_level }}-gpu=${{ inputs.gpu }}-cuvs=${{ inputs.cuvs }}-rocm=${{ inputs.rocm }}
         path: test-results
+    - name: Check installed packages channel
+      shell: bash
+      run: |
+        # Shows that all installed packages are from conda-forge.
+        conda list --show-channel-urls
diff --git a/.github/actions/build_conda/action.yml b/.github/actions/build_conda/action.yml
@@ -30,16 +30,22 @@ runs:
       uses: conda-incubator/setup-miniconda@v3
       with:
         python-version: '3.11'
-        miniconda-version:  latest
+        miniforge-version: latest # ensures conda-forge channel is used.
+        channels: conda-forge
+        conda-remove-defaults: 'true'
+        # Set to runner.arch=aarch64 if we're on arm64 because
+        # there's no miniforge ARM64 package, just aarch64.
+        # They are the same thing, just named differently.
+        # However there is an ARM64 for macOS, so exclude that.
+        architecture: ${{ (runner.arch == 'ARM64' && runner.os != 'macOS') && 'aarch64' || runner.arch }}
     - name: Install conda build tools
       shell: ${{ steps.choose_shell.outputs.shell }}
       run: |
+        # Ensure starting packages are from conda-forge.
+        conda list --show-channel-urls
         conda install -y -q "conda!=24.11.0"
         conda install -y -q "conda-build!=24.11.0"
-    - name: Fix CI failure
-      shell: ${{ steps.choose_shell.outputs.shell }}
-      if: runner.os != 'Windows'
-      run: conda remove conda-anaconda-telemetry
+        conda list --show-channel-urls
     - name: Enable anaconda uploads
       if: inputs.label != ''
       shell: ${{ steps.choose_shell.outputs.shell }}
@@ -94,3 +100,8 @@ runs:
       run: |
         conda build faiss-gpu-cuvs --variants '{ "cudatoolkit": "${{ inputs.cuda }}" }' \
             --user pytorch --label ${{ inputs.label }} -c pytorch -c rapidsai -c rapidsai-nightly -c conda-forge -c nvidia
+    - name: Check installed packages channel
+      shell: ${{ steps.choose_shell.outputs.shell }}
+      run: |
+        # Shows that all installed packages are from conda-forge.
+        conda list --show-channel-urls
diff --git a/conda/faiss-gpu/meta.yaml b/conda/faiss-gpu/meta.yaml
@@ -50,14 +50,16 @@ outputs:
         - sysroot_linux-64 =2.17 # [linux64]
         - llvm-openmp  # [osx]
         - cmake >=3.24.0
-        - make =4.2 # [not win]
-        - mkl-devel =2023  # [x86_64]
+        - make =4.2 # [not win and not (osx and arm64)]
+        - make =4.4 # [osx and arm64]
+        - mkl-devel =2023.0  # [x86_64]
         - cuda-toolkit {{ cudatoolkit }}
+        - gcc_linux-64 =11.2  # [cudatoolkit == '11.4.4']
       host:
-        - mkl =2023  # [x86_64]
+        - mkl =2023.0  # [x86_64]
         - openblas =0.3 # [not x86_64]
       run:
-        - mkl =2023  # [x86_64]
+        - mkl =2023.0  # [x86_64]
         - openblas =0.3 # [not x86_64]
         - cuda-cudart {{ cuda_constraints }}
         - libcublas {{ libcublas_constraints }}
@@ -83,11 +85,14 @@ outputs:
         - sysroot_linux-64 =2.17 # [linux64]
         - swig =4.0
         - cmake >=3.24.0
-        - make =4.2 # [not win]
+        - make =4.2 # [not win and not (osx and arm64)]
+        - make =4.4 # [osx and arm64]
+        - _openmp_mutex =4.5=2_kmp_llvm  # [x86_64 and not win]
         - cuda-toolkit {{ cudatoolkit }}
       host:
         - python {{ python }}
         - numpy >=1.19,<2
+        - _openmp_mutex =4.5=2_kmp_llvm  # [x86_64 and not win]
         - {{ pin_subpackage('libfaiss', exact=True) }}
       run:
         - python {{ python }}

diff --git a/conda/faiss/meta.yaml b/conda/faiss/meta.yaml
@@ -31,22 +31,53 @@ outputs:
     script: build-lib-arm64.sh  # [not x86_64]
     script: build-lib.bat  # [win]
     build:
-      string: "h{{ PKG_HASH }}_{{ number }}_cpu{{ suffix }}"
+      string: "py{{ PY_VER }}_h{{ PKG_HASH }}_{{ number }}_cpu{{ suffix }}"
       run_exports:
         - {{ pin_compatible('libfaiss', exact=True) }}
     requirements:
       build:
+        - python {{ python }}
         - {{ compiler('cxx') }}
         - sysroot_linux-64 =2.17 # [linux64]
-        - llvm-openmp  # [osx]
+        - llvm-openmp  # [osx or linux64]
         - cmake >=3.24.0
-        - make =4.2 # [not win]
-        - mkl-devel =2023  # [x86_64]
+        - make =4.2 # [not win and not (osx and arm64)]
+        - make =4.4 # [osx and arm64]
+        {% if PY_VER == '3.9' or PY_VER == '3.10' or PY_VER == '3.11' %}
+        - mkl-devel =2023.0  # [x86_64]
+        - liblief =0.12.3  # [not win]
+        - python_abi <3.12
+        {% elif PY_VER == '3.12' %}
+        - mkl-devel >=2023.2.0  # [x86_64 and not win]
+        - mkl-devel =2023.1.0  # [x86_64 and win]
+        - liblief =0.15.1  # [not win]
+        - python_abi =3.12
+        {% endif %}
       host:
-        - mkl =2023  # [x86_64]
+        - python {{ python }}
+        {% if PY_VER == '3.9' or PY_VER == '3.10' or PY_VER == '3.11' %}
+        - mkl =2023.0  # [x86_64]
+        - liblief =0.12.3  # [not win]
+        - python_abi <3.12
+        {% elif PY_VER == '3.12' %}
+        - mkl >=2023.2.0  # [x86_64 and not win]
+        - mkl =2023.1.0  # [x86_64 and win]
+        - liblief =0.15.1  # [not win]
+        - python_abi =3.12
+        {% endif %}
         - openblas =0.3 # [not x86_64]
       run:
-        - mkl =2023  # [x86_64]
+        - python {{ python }}
+        {% if PY_VER == '3.9' or PY_VER == '3.10' or PY_VER == '3.11' %}
+        - mkl =2023.0  # [x86_64]
+        - liblief =0.12.3  # [not win]
+        - python_abi <3.12
+        {% elif PY_VER == '3.12' %}
+        - mkl >=2023.2.0  # [x86_64 and not win]
+        - mkl =2023.1.0  # [x86_64 and win]
+        - liblief =0.15.1  # [not win]
+        - python_abi =3.12
+        {% endif %}
         - openblas =0.3 # [not x86_64]
     test:
       requires:
@@ -66,25 +97,61 @@ outputs:
       string: "py{{ PY_VER }}_h{{ PKG_HASH }}_{{ number }}_cpu{{ suffix }}"
     requirements:
       build:
+        - python {{ python }}
         - {{ compiler('cxx') }}
         - sysroot_linux-64 =2.17 # [linux64]
         - swig =4.0
         - cmake >=3.24.0
-        - make =4.2 # [not win]
+        - make =4.2 # [not win and not (osx and arm64)]
+        - make =4.4 # [osx and arm64]
+        - _openmp_mutex =4.5=2_kmp_llvm  # [x86_64 and not win]
+        {% if PY_VER == '3.9' or PY_VER == '3.10' or PY_VER == '3.11' %}
+        - mkl =2023.0  # [x86_64]
+        - python_abi <3.12
+        {% elif PY_VER == '3.12' %}
+        - mkl >=2023.2.0  # [x86_64 and not win]
+        - mkl =2023.1.0  # [x86_64 and win]
+        - python_abi =3.12
+        {% endif %}
       host:
         - python {{ python }}
         - numpy >=1.19,<2
         - {{ pin_subpackage('libfaiss', exact=True) }}
+        - _openmp_mutex =4.5=2_kmp_llvm  # [x86_64 and not win]
+        {% if PY_VER == '3.9' or PY_VER == '3.10' or PY_VER == '3.11' %}
+        - mkl =2023.0  # [x86_64]
+        - python_abi <3.12
+        {% elif PY_VER == '3.12' %}
+        - mkl >=2023.2.0  # [x86_64 and not win]
+        - mkl =2023.1.0  # [x86_64 and win]
+        - python_abi =3.12
+        {% endif %}
       run:
         - python {{ python }}
         - numpy >=1.19,<2
         - packaging
         - {{ pin_subpackage('libfaiss', exact=True) }}
+        {% if PY_VER == '3.9' or PY_VER == '3.10' or PY_VER == '3.11' %}
+        - mkl =2023.0  # [x86_64]
+        - python_abi <3.12
+        {% elif PY_VER == '3.12' %}
+        - mkl >=2023.2.0  # [x86_64 and not win]
+        - mkl =2023.1.0  # [x86_64 and win]
+        - python_abi =3.12
+        {% endif %}
     test:
       requires:
         - numpy >=1.19,<2
         - scipy
         - pytorch <2.5
+        {% if PY_VER == '3.9' or PY_VER == '3.10' or PY_VER == '3.11' %}
+        - mkl =2023.0  # [x86_64]
+        - python_abi <3.12
+        {% elif PY_VER == '3.12' %}
+        - mkl >=2023.2.0  # [x86_64 and not win]
+        - mkl =2023.1.0  # [x86_64 and win]
+        - python_abi =3.12
+        {% endif %}
       commands:
         - python -X faulthandler -m unittest discover -v -s tests/ -p "test_*"
         - python -X faulthandler -m unittest discover -v -s tests/ -p "torch_*"

diff --git a/faiss/gpu/test/CMakeLists.txt b/faiss/gpu/test/CMakeLists.txt
@@ -45,6 +45,7 @@ faiss_gpu_test(TestGpuIndexBinaryFlat.cpp)
 faiss_gpu_test(TestGpuMemoryException.cpp)
 faiss_gpu_test(TestGpuIndexIVFPQ.cpp)
 faiss_gpu_test(TestGpuIndexIVFScalarQuantizer.cpp)
+faiss_gpu_test(TestGpuResidualQuantizer.cpp)
 faiss_gpu_test(TestGpuDistance.${GPU_EXT_PREFIX})
 faiss_gpu_test(TestGpuSelect.${GPU_EXT_PREFIX})
 if(FAISS_ENABLE_CUVS)

diff --git a/faiss/gpu/test/TestGpuResidualQuantizer.cpp b/faiss/gpu/test/TestGpuResidualQuantizer.cpp
@@ -0,0 +1,70 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#include <faiss/IndexFlat.h>
+#include <faiss/gpu/GpuCloner.h>
+#include <faiss/gpu/GpuIndexFlat.h>
+#include <faiss/gpu/StandardGpuResources.h>
+#include <faiss/gpu/test/TestUtils.h>
+#include <faiss/impl/ResidualQuantizer.h>
+#include <gtest/gtest.h>
+
+using namespace ::testing;
+
+float eval_codec(faiss::ResidualQuantizer* q, int nb, float* xb) {
+    // Compute codes
+    uint8_t* codes = new uint8_t[q->code_size * nb];
+    std::cout << "code size: " << q->code_size << std::endl;
+    q->compute_codes(xb, codes, nb);
+    // Decode codes
+    float* decoded = new float[nb * q->d];
+    q->decode(codes, decoded, nb);
+    // Compute reconstruction error
+    float err = 0.0f;
+    for (int i = 0; i < nb; i++) {
+        for (int j = 0; j < q->d; j++) {
+            float diff = xb[i * q->d + j] - decoded[i * q->d + j];
+            err = err + (diff * diff);
+        }
+    }
+    delete[] codes;
+    delete[] decoded;
+    return err;
+}
+
+TEST(TestGpuResidualQuantizer, TestNcall) {
+    int d = 32;
+    int nt = 3000;
+    int nb = 1000;
+    // Assuming get_dataset_2 is a function that returns xt and xb
+    std::vector<float> xt = faiss::gpu::randVecs(nt, d);
+    std::vector<float> xb = faiss::gpu::randVecs(nb, d);
+    faiss::ResidualQuantizer rq0(d, 4, 6);
+    rq0.train(nt, xt.data());
+    float err_rq0 = eval_codec(&rq0, nb, xb.data());
+    faiss::ResidualQuantizer rq1(d, 4, 6);
+    faiss::gpu::GpuProgressiveDimIndexFactory fac(1);
+    rq1.assign_index_factory = &fac;
+    rq1.train(nt, xt.data());
+    ASSERT_GT(fac.ncall, 0);
+    int ncall_train = fac.ncall;
+    float err_rq1 = eval_codec(&rq1, nb, xb.data());
+    ASSERT_GT(fac.ncall, ncall_train);
+    std::cout << "Error RQ0: " << err_rq0 << ", Error RQ1: " << err_rq1
+              << std::endl;
+    ASSERT_TRUE(0.9 * err_rq0 < err_rq1);
+    ASSERT_TRUE(err_rq1 < 1.1 * err_rq0);
+}
+
+int main(int argc, char** argv) {
+    testing::InitGoogleTest(&argc, argv);
+
+    // just run with a fixed test seed
+    faiss::gpu::setTestSeed(100);
+
+    return RUN_ALL_TESTS();
+}
diff --git a/faiss/gpu/test/test_gpu_basics.py b/faiss/gpu/test/test_gpu_basics.py
@@ -428,6 +428,15 @@ def eval_codec(q, xb):
 
 class TestResidualQuantizer(unittest.TestCase):
 
+    # This test is disabled due to memory corruption in some dependency.
+    # It only happens in CUDA 11.4.4 after switching from  defaults
+    # to conda-forge for dependencies.
+    # GpuProgressiveDimIndexFactory is partially overwritten, and ncall
+    # ends up with garbage data when checking it in Python. However,
+    # the C++ side prints the right values. This is likely a compiler bug.
+    # This test is left in the codebase for now but skipped so that we
+    # know there is a problem with it.
+    @unittest.skip("Skipped due to ncall memory corruption.")
     def test_with_gpu(self):
         """ check that we get the same results with a GPU quantizer and a CPU quantizer """
         d = 32

diff --git a/tests/test_contrib.py b/tests/test_contrib.py
@@ -12,6 +12,7 @@
 
 import faiss
 import numpy as np
+import sys
 
 from common_faiss_tests import get_dataset_2
 
@@ -392,6 +393,12 @@ def test_float(self):
             l0, l1 = lims[q], lims[q + 1]
             self.assertTrue(set(I[q]) <= set(IR[l0:l1]))
 
+    @unittest.skipIf(
+        platform.system() == 'Windows'
+        and sys.version_info[0] == 3
+        and sys.version_info[1] == 12,
+        'test_binary hangs for Windows on Python 3.12.'
+    )
     def test_binary(self):
         ds = datasets.SyntheticDataset(128, 2000, 2000, 200)