SNMG ANN build with OpenMP nested parallelism

[viclafargue]

The SNMG ANN build spawns threads with #pragma omp parallel for, but CAGRA build also uses OpenMP internally. With nested parallelism disabled by default, CAGRA's internal parallel regions ran serially on a single thread, causing severe performance degradation.

• Enable nested parallelism with omp_set_nested(1)
• Limit outer loop to num_ranks threads (one per GPU)
• Inside each rank thread, allocate threads_per_rank for CAGRA's internal parallelism
• Restore original thread count after parallel region

[tarang-jain]

this restoration is a bit confusing. the pool of threads is set inside the for loop above, but its restored after the for loop? Can you verify if this logic is correct?

[viclafargue]

I agree. This is actually not necessary. My worry was that the main thread could be one of the threads in the loop. But, after checking it appears like OpenMP is designed in such a way that omp_set_num_threads calls inside of a parallel region only affect how many threads each thread can use in nested parallel region. The main thread is unaffected, and its number of threads is actually preserved. I made sure that each thread sees its number of thread restored back to the global maximal number of threads within the parallel region.

[viclafargue]

I also added the same OpenMP usage for the extend function.

[viclafargue]

/merge