Fix reverse_index_factory/get_code_size round-trip for HNSW, IMI and IVFPQR (#5564)

Summary:
Three defects in `contrib/factory_tools.py` where a reversed factory string rebuilds a different index or fails to parse. `test_get_code_size_hnsw_roundtrip` already pins this invariant for flat HNSW; these cases escape it.

1. **HNSW storage dropped.** The `IndexHNSW` branch never read `index.storage`, so `HNSW32,SQ8` and `HNSW32,PQ16x8` both reversed to `"HNSW32"` and rebuilt as `IndexHNSWFlat`. `get_code_size` then applied the flat `d*4` formula: at d=128 it reported 768 bytes/vector for an index whose storage is 128 plus the 256-byte graph estimate, so 384.

2. **IMI regex fields transposed.** `index_factory.cpp:295` parses `IMI2x([0-9]+)`, but the regex read `IMI\d+x2`. Every real IMI key (`IMI2x10,PQ16`, and the ones in `demos/demo_auto_tune.py`) raised `RuntimeError("cannot parse")`.

3. **IndexIVFPQR matched by the IndexIVFPQ check**, dropping the refinement stage: `IVF64,PQ8+16` reversed to `IVF64,PQ8x8`, 8 bytes instead of 24.

Consumers are `benchs/bench_all_ivf/parse_bench_all_ivf.py` (imports `get_code_size` as `unitsize`) and `benchs/bench_fw/index.py`.

Flat storage stays implicit as the bare `HNSW<M>` string, and a null `storage` pointer still returns `HNSW<M>` instead of raising.

**Behavior change worth flagging:** `IndexHNSW2Level` storage is `Index2Layer`, which `reverse_index_factory` does not handle, so it now raises `NotImplementedError` instead of silently returning a wrong `"HNSW32"`. It is not reachable from `index_factory`.

Pull Request resolved: #5564

Test Plan:
Four tests added to `TestFactoryTools`. All four fail on unpatched `factory_tools.py` and pass with it:

```
# pristine contrib/factory_tools.py, new tests present
4 failed, 7 passed
FAILED test_get_code_size_hnsw_non_flat_storage
FAILED test_get_code_size_imi
FAILED test_hnsw_storage_reverse_index_factory
FAILED test_ivfpqr_reverse_index_factory

# with the fix
11 passed
```

Full `tests/test_contrib.py tests/test_factory.py`, pristine vs branch:

```
pristine: 2 failed, 88 passed
branch:   2 failed, 92 passed
```

Failure sets are identical. The two are `TestComputeGT::test_compute_GT_gpu` and `test_compute_GT_ip_gpu`, pre-existing and unrelated (`AttributeError: module 'faiss' has no attribute 'GpuMultipleClonerOptions'` on a CPU-only build).

An A/B sweep over 40 factory strings changes exactly the 7 expected lines and nothing else:

```
- HNSW32,SQ8      IndexHNSWSQ   HNSW32        768
+ HNSW32,SQ8      IndexHNSWSQ   HNSW32,SQ8    384
- HNSW32,SQfp16   IndexHNSWSQ   HNSW32        768
+ HNSW32,SQfp16   IndexHNSWSQ   HNSW32,SQfp16 512
- HNSW32,PQ16x8   IndexHNSWPQ   HNSW32        768
+ HNSW32,PQ16x8   IndexHNSWPQ   HNSW32,PQ16x8 272
- IVF64,PQ8+16    IndexIVFPQR   IVF64,PQ8x8   8
+ IVF64,PQ8+16    IndexIVFPQR   IVF64,PQ8+16  24
- IMI2x5,Flat     IndexIVFFlat  IMI2x5,Flat   <RuntimeError>
+ IMI2x5,Flat     IndexIVFFlat  IMI2x5,Flat   512
- IMI2x5,PQ8      IndexIVFPQ    IMI2x5,PQ8x8  <RuntimeError>
+ IMI2x5,PQ8      IndexIVFPQ    IMI2x5,PQ8x8  8
- IMI2x8,PQ8x8    IndexIVFPQ    IMI2x8,PQ8x8  <RuntimeError>
+ IMI2x8,PQ8x8    IndexIVFPQ    IMI2x8,PQ8x8  8
```

flake8 at the 80-char limit from CONTRIBUTING reports the same 12 pre-existing findings before and after, none in the changed lines.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Reviewed By: mnorris11

Differential Revision: D118661065

Pulled By: trang-nm-nguyen

fbshipit-source-id: b50a9522361e3fd14f276b36196d3c728567b980

Author: Kayvan-Zahiri

contrib/factory_tools.py

@@ -25,7 +25,7 @@ def get_code_size(d, indexkey):
     if mo:
         return get_code_size(d, mo.group(1))
 
-    mo = re.match("IMI\\d+x2,(.*)$", indexkey)
+    mo = re.match("IMI\\d+x\\d+,(.*)$", indexkey)
     if mo:
         return get_code_size(d, mo.group(1))
 
@@ -50,6 +50,11 @@ def get_code_size(d, indexkey):
         M = int(mo.group(1))
         return d * 4 + M * 2 * 4  # roughly
 
+    mo = re.match("HNSW(\\d+),(.*)$", indexkey)
+    if mo:
+        M = int(mo.group(1))
+        return get_code_size(d, mo.group(2)) + M * 2 * 4  # roughly
+
     if indexkey == "SQ8":
         return d
     elif indexkey == "SQ4":
@@ -127,6 +132,9 @@ def reverse_index_factory(index):
             return prefix + ",Flat"
         if isinstance(index, faiss.IndexIVFScalarQuantizer):
             return prefix + "," + sq_names[index.sq.qtype]
+        # IndexIVFPQR subclasses IndexIVFPQ, so it must be checked first.
+        if isinstance(index, faiss.IndexIVFPQR):
+            return prefix + f",PQ{index.pq.M}+{index.refine_pq.M}"
         if isinstance(index, faiss.IndexIVFPQ):
             return prefix + f",PQ{index.pq.M}x{index.pq.nbits}"
         if isinstance(index, faiss.IndexIVFPQFastScan):
@@ -152,7 +160,12 @@ def reverse_index_factory(index):
         return f"{prefix},{reverse_index_factory(index.index)}"
 
     elif isinstance(index, faiss.IndexHNSW):
-        return f"HNSW{get_hnsw_M(index)}"
+        prefix = f"HNSW{get_hnsw_M(index)}"
+        storage = faiss.downcast_index(index.storage)
+        # flat storage is implicit in the bare HNSW<M> factory string
+        if storage is None or isinstance(storage, faiss.IndexFlat):
+            return prefix
+        return f"{prefix},{reverse_index_factory(storage)}"
 
     elif isinstance(index, faiss.IndexRefine):
         return (

tests/test_contrib.py

@@ -1021,6 +1021,56 @@ def test_get_code_size_hnsw_roundtrip(self):
         code_size = factory_tools.get_code_size(d, factory_str)
         self.assertEqual(code_size, d * 4 + 16 * 2 * 4)
 
+    def test_hnsw_storage_reverse_index_factory(self):
+        d = 128
+        # flat storage stays implicit, everything else must be spelled out
+        cases = {
+            "HNSW32,Flat": "HNSW32",
+            "HNSW32,SQ8": "HNSW32,SQ8",
+            "HNSW32,SQfp16": "HNSW32,SQfp16",
+            "HNSW32,PQ16x8": "HNSW32,PQ16x8",
+        }
+        for key, expected in cases.items():
+            index = faiss.index_factory(d, key)
+            factory_str = factory_tools.reverse_index_factory(index)
+            self.assertEqual(factory_str, expected)
+            # the string must rebuild the same index class
+            rebuilt = faiss.index_factory(d, factory_str)
+            self.assertEqual(type(rebuilt), type(index))
+
+    def test_get_code_size_hnsw_non_flat_storage(self):
+        d = 128
+        graph = 32 * 2 * 4
+        self.assertEqual(
+            factory_tools.get_code_size(d, "HNSW32,SQ8"), d + graph
+        )
+        self.assertEqual(
+            factory_tools.get_code_size(d, "HNSW32,SQfp16"), d * 2 + graph
+        )
+        self.assertEqual(
+            factory_tools.get_code_size(d, "HNSW32,PQ16x8"), 16 + graph
+        )
+
+    def test_get_code_size_imi(self):
+        # the coarse quantizer adds no per-vector bytes
+        d = 64
+        self.assertEqual(factory_tools.get_code_size(d, "IMI2x10,PQ16"), 16)
+        self.assertEqual(factory_tools.get_code_size(d, "IMI2x5,Flat"), d * 4)
+        # round-trip: get_code_size must parse what reverse emits
+        index = faiss.index_factory(d, "IMI2x5,PQ8")
+        factory_str = factory_tools.reverse_index_factory(index)
+        self.assertEqual(factory_str, "IMI2x5,PQ8x8")
+        self.assertEqual(factory_tools.get_code_size(d, factory_str), 8)
+
+    def test_ivfpqr_reverse_index_factory(self):
+        d = 128
+        index = faiss.index_factory(d, "IVF64,PQ8+16")
+        factory_str = factory_tools.reverse_index_factory(index)
+        self.assertEqual(factory_str, "IVF64,PQ8+16")
+        rebuilt = faiss.index_factory(d, factory_str)
+        self.assertEqual(type(rebuilt), type(index))
+        self.assertEqual(factory_tools.get_code_size(d, factory_str), 8 + 16)
+
     def test_rabitq_reverse_index_factory(self):
         d = 64
         quantizer = faiss.IndexFlatL2(d)