From e770054b7f60df7bc36eff5dcf1347ee0baf7fd8 Mon Sep 17 00:00:00 2001
From: Wenkai Du <43822138+wenkaidu@users.noreply.github.com>
Date: Fri, 16 Aug 2024 06:16:00 +0800
Subject: [PATCH] Fix model matching with PXN enable (#1295)

[ROCm/rccl commit: eff56735b05194ad06205ad3f91452014bc6f792]
---
 projects/rccl/src/graph/rome_models.cc | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/projects/rccl/src/graph/rome_models.cc b/projects/rccl/src/graph/rome_models.cc
index 19556746f8..5b70ee1d40 100644
--- a/projects/rccl/src/graph/rome_models.cc
+++ b/projects/rccl/src/graph/rome_models.cc
@@ -1475,6 +1475,8 @@ static bool permuteNetIds(int *n, int *g, int s, int last, struct rcclRomeModel*
     // match gdr level
     for (i = 0; i < ref->nNics; i++) {
       for (j = 0; j < ref->nGpus; j++) {
+        // enabling PXN override paths over PHB and SYS
+        if (topo->gdrLevel[n[i]*ref->nGpus+g[j]] == PATH_PXN) continue;
         if (ref->gdrLevel[i*ref->nGpus+j] != topo->gdrLevel[n[i]*ref->nGpus+g[j]]) break;
       }
       if (j < ref->nGpus) break;
@@ -1565,8 +1567,11 @@ ncclResult_t parseRome4P2H(struct ncclTopoSystem* system, struct ncclTopoGraph*
           if (!ignore_numa && romeTopoModels[i].nicNuma[j] != romeTopo.nicNuma[k]) continue;
           int g;
           // check GDR
-          for (g = 0; g < ngpus; g++)
+          for (g = 0; g < ngpus; g++) {
+            // enabling PXN override paths over PHB and SYS
+            if (romeTopo.gdrLevel[k*ngpus+g] == PATH_PXN) continue;
             if (romeTopoModels[i].gdrLevel[j*ngpus+g] != romeTopo.gdrLevel[k*ngpus+g]) break;
+          }
           if (g >= ngpus) break;
         }
         if (k < nnets) {