Fix memory allocation bug to use the correct hipDevice (#308)

- Not using hipSetDevice before allocating memory can use unintended
deviceIdx when executing fabric-handle based transfers
- Reset numa_set_preferred(-1) before ERR_FATAL early return in the
  non-POD_COMM_ENABLED path; without this the NUMA policy stays dirty
  for subsequent CPU allocations in the same process
- Use memDevice.memIndex directly in the top-level hipSetDevice call
  instead of deviceIdx, which is NUMA-remapped for CPU types only;
  documents that the MEM_CPU_CLOSEST remapping does not apply to GPU
- Remove now-redundant hipSetDevice inside the POD_COMM GPU memHandle
  branch; device was already set at the top of AllocateMemory
- Guard CollectTopology GPU agent probe loop with hipSetDevice(i) so
  each AllocateMemory call targets the correct device

---------

Co-authored-by: Claude <claude@anthropic.com>

Author: nileshnegi

src/header/TransferBench.hpp

@@ -1496,6 +1496,15 @@ namespace {
       deviceIdx = GetClosestCpuNumaToGpu(memDevice.memIndex);
     }
 
+    if (IsCpuMemType(memType)) {
+      // Set NUMA policy prior to call to hipHostMalloc
+      numa_set_preferred(deviceIdx);
+    } else if (IsGpuMemType(memType)) {
+      // Switch to the appropriate GPU
+      // IMP: if the remapping above changes, remember to modify this!
+      ERR_CHECK(hipSetDevice(deviceIdx));
+    }
+
     // If memHandle is provided, allocate sharable memory
     if (memHandle != NULL) {
 #ifdef POD_COMM_ENABLED
@@ -1533,13 +1542,14 @@ namespace {
         memset(*memPtr, 0, roundedUpBytes);
         // Check that the allocated pages are actually on the correct NUMA node
         ERR_CHECK(CheckPages((char*)*memPtr, roundedUpBytes, deviceIdx));
+        numa_set_preferred(-1);
       } else if (IsGpuMemType(memType)) {
-        ERR_CHECK(hipSetDevice(memDevice.memIndex));
         ERR_CHECK(hipMemset(*memPtr, 0, numBytes));
         ERR_CHECK(hipDeviceSynchronize());
       }
       return ERR_NONE;
 #else
+      if (IsCpuMemType(memType)) numa_set_preferred(-1);
       return {ERR_FATAL, "Unable to allocate sharable memory if not compiled with pod communication support"};
 #endif
     } else {
@@ -1548,9 +1558,6 @@ namespace {
 
     if (IsCpuMemType(memType)) {
 
-      // Set NUMA policy prior to call to hipHostMalloc
-      numa_set_preferred(deviceIdx);
-
       // Allocate host-pinned memory (should respect NUMA mem policy)
       int flags = 0;
 #if !defined (__NVCC__)
@@ -1591,8 +1598,6 @@ namespace {
       // Reset to default numa mem policy
       numa_set_preferred(-1);
     } else if (IsGpuMemType(memType)) {
-      // Switch to the appropriate GPU
-      ERR_CHECK(hipSetDevice(memDevice.memIndex));
 
       if (memType == MEM_GPU) {
         // Allocate GPU memory on appropriate device