diff --git a/faiss/gpu/test/test_lut_dtype_binding.py b/faiss/gpu/test/test_lut_dtype_binding.py new file mode 100644 index 000000000..079ff53cb --- /dev/null +++ b/faiss/gpu/test/test_lut_dtype_binding.py @@ -0,0 +1,60 @@ +# 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. + +import unittest + +import faiss + + +@unittest.skipIf( + "CUVS" not in faiss.get_compile_options(), + "only if cuVS is compiled in") +class TestIVFPQSearchCagraConfigDtypes(unittest.TestCase): + """Regression tests for the SWIG int typemap on cudaDataType_t. + + Before the typemap was added, IVFPQSearchCagraConfig.lut_dtype and + .internal_distance_dtype were exposed as SwigPyObject pointers that + Python users could neither construct nor assign to, and the + CUDA_R_* enum values were not exported. These tests pin the + behavior that they are now plain ints settable from Python. + """ + + def test_constants_exported(self): + self.assertEqual(faiss.CUDA_R_32F, 0) + self.assertEqual(faiss.CUDA_R_64F, 1) + self.assertEqual(faiss.CUDA_R_16F, 2) + self.assertEqual(faiss.CUDA_R_8I, 3) + self.assertEqual(faiss.CUDA_R_8U, 8) + + def test_default_lut_dtype_is_fp32(self): + c = faiss.IVFPQSearchCagraConfig() + self.assertEqual(c.lut_dtype, faiss.CUDA_R_32F) + self.assertEqual(c.internal_distance_dtype, faiss.CUDA_R_32F) + + def test_set_lut_dtype_via_constants(self): + c = faiss.IVFPQSearchCagraConfig() + for value in ( + faiss.CUDA_R_16F, + faiss.CUDA_R_8U, + faiss.CUDA_R_32F, + ): + c.lut_dtype = value + self.assertEqual(c.lut_dtype, value) + + def test_set_lut_dtype_via_raw_int(self): + c = faiss.IVFPQSearchCagraConfig() + c.lut_dtype = 2 # CUDA_R_16F + self.assertEqual(c.lut_dtype, 2) + + def test_dtype_fields_are_independent(self): + c = faiss.IVFPQSearchCagraConfig() + c.lut_dtype = faiss.CUDA_R_8U + c.internal_distance_dtype = faiss.CUDA_R_16F + self.assertEqual(c.lut_dtype, faiss.CUDA_R_8U) + self.assertEqual(c.internal_distance_dtype, faiss.CUDA_R_16F) + + +if __name__ == "__main__": + unittest.main() diff --git a/faiss/python/swigfaiss.swig b/faiss/python/swigfaiss.swig index 74c9cea2d..12be4c343 100644 --- a/faiss/python/swigfaiss.swig +++ b/faiss/python/swigfaiss.swig @@ -727,6 +727,25 @@ void gpu_sync_all_devices() %include %include #ifdef FAISS_ENABLE_CUVS +// IVFPQSearchCagraConfig.lut_dtype and .internal_distance_dtype are typed as +// cudaDataType_t (a C enum from ). Without an int typemap +// SWIG generates accessors that read/write a cudaDataType_t* SwigPyObject, +// which Python users cannot construct, so the fields are effectively +// read-only from Python and the documented values (CUDA_R_16F, CUDA_R_8U) +// are not exported. Treat cudaDataType_t as int for SWIG and expose the +// relevant enum values so: +// +// c = faiss.IVFPQSearchCagraConfig() +// c.lut_dtype = faiss.CUDA_R_16F +// +// works as expected. Values match cudaDataType_t in . +typedef int cudaDataType_t; +%apply int { cudaDataType_t }; +%constant int CUDA_R_32F = 0; +%constant int CUDA_R_64F = 1; +%constant int CUDA_R_16F = 2; +%constant int CUDA_R_8I = 3; +%constant int CUDA_R_8U = 8; %include %include #endif