prov/opx: Use appropriate memory monitor when registering TID regions.

Signed-off-by: Ben Lynam <[email protected]>
ofiwg · Jul 1, 2024 · 076754a · 076754a
1 parent 28a67f1
commit 076754a
Showing 1 changed file with 8 additions and 3 deletions.
diff --git a/prov/opx/src/fi_opx_tid_cache.c b/prov/opx/src/fi_opx_tid_cache.c
@@ -898,9 +898,7 @@ int opx_tid_cache_crte(struct ofi_mr_cache *cache,
 {
 	/*struct ofi_mr_entry *cur;*/
 	int ret;
-	// TODO - Original behavior is struct ofi_mem_monitor *monitor = cache->monitors[info->iface];
-	// Hardcoding to FI_HMEM_SYSTEM
-	struct ofi_mem_monitor *monitor = cache->monitors[FI_HMEM_SYSTEM];
+	struct ofi_mem_monitor *monitor = cache->monitors[info->iface];
 
 	/* Assert precondition that the lock is held with a trylock assert */
 	assert(pthread_mutex_trylock(&mm_lock) == EBUSY);
@@ -1289,6 +1287,13 @@ int opx_tid_cache_setup(struct ofi_mr_cache **cache,
 
 	struct ofi_mem_monitor *memory_monitors[OFI_HMEM_MAX] = {
 		[FI_HMEM_SYSTEM] = default_monitor,
+		[FI_HMEM_CUDA] = default_cuda_monitor,
+		[FI_HMEM_ROCR] = default_rocr_monitor,
+		[FI_HMEM_ZE] = default_ze_monitor,
+
+		/* Not supported in OPX */
+		[FI_HMEM_NEURON] = NULL,
+		[FI_HMEM_SYNAPSEAI] = NULL
 	};
 	int err;