From 18d6484c67cc87eb409c474206e8cbb06bccb343 Mon Sep 17 00:00:00 2001
From: "Bill(Shuzhou) Liu" <shuzhou.liu@amd.com>
Date: Mon, 6 Nov 2023 13:44:31 -0600
Subject: [PATCH] Sort GPU index using BDF

Sort GPU index based on BDF.

Change-Id: I921fc72e534665bd14b9a45a747f72b24f5ba641
---
 rocm_smi/src/rocm_smi_main.cc | 26 ++++++++++++++++++++++++--
 1 file changed, 24 insertions(+), 2 deletions(-)

diff --git a/rocm_smi/src/rocm_smi_main.cc b/rocm_smi/src/rocm_smi_main.cc
index c29916bedc..98fae2a663 100755
--- a/rocm_smi/src/rocm_smi_main.cc
+++ b/rocm_smi/src/rocm_smi_main.cc
@@ -52,6 +52,7 @@
 #include <fstream>
 #include <functional>
 #include <iostream>
+#include <algorithm>
 #include <memory>
 #include <set>
 #include <sstream>
@@ -393,6 +394,29 @@ RocmSMI::Initialize(uint64_t flags) {
             "Failed to initialize rocm_smi library (amdgpu node discovery).");
   }
 
+  std::shared_ptr<amd::smi::Device> dev;
+  // Sort index based on the BDF, collect BDF id firstly.
+  std::vector<std::pair<uint64_t, std::shared_ptr<amd::smi::Device>>> dv_to_id;
+  dv_to_id.reserve(devices_.size());
+  for (uint32_t dv_ind = 0; dv_ind < devices_.size(); ++dv_ind) {
+      dev = devices_[dv_ind];
+      uint64_t bdfid = dev->bdfid();
+      dv_to_id.push_back({bdfid, dev});
+    }
+  ss << __PRETTY_FUNCTION__ << " Sort index based on BDF.";
+  LOG_DEBUG(ss);
+
+  // Stable sort to keep the order if bdf is equal.
+  std::stable_sort(dv_to_id.begin(), dv_to_id.end(), []
+  (const std::pair<uint64_t, std::shared_ptr<amd::smi::Device>>& p1,
+    const std::pair<uint64_t, std::shared_ptr<amd::smi::Device>>& p2) {
+        return p1.first < p2.first;
+  });
+  devices_.clear();
+  for (uint32_t dv_ind = 0; dv_ind < dv_to_id.size(); ++dv_ind) {
+    devices_.push_back(dv_to_id[dv_ind].second);
+  }
+
   std::map<uint64_t, std::shared_ptr<KFDNode>> tmp_map;
   i_ret = DiscoverKFDNodes(&tmp_map);
   if (i_ret != 0) {
@@ -411,8 +435,6 @@ RocmSMI::Initialize(uint64_t flags) {
   for (it = io_link_map_tmp.begin(); it != io_link_map_tmp.end(); it++)
     io_link_map_[it->first] = it->second;
 
-  std::shared_ptr<amd::smi::Device> dev;
-
   // Remove any drm nodes that don't have  a corresponding readable kfd node.
   // kfd nodes will not be added if their properties file is not readable.
   auto dev_iter = devices_.begin();