Add MPI rank support for PyTorch profiler output files

Co-authored-by: caic99 <78061359+caic99@users.noreply.github.com>

Author: Copilot

PYTORCH_PROFILER_INTEGRATION.md

@@ -14,7 +14,11 @@ export DP_PYTORCH_PROFILER_OUTPUT_DIR=./profiler_results
 
 3. Check for profiler output in the specified directory:
 ```bash
+# For single-rank or non-MPI usage
 ls -la ./profiler_results/pytorch_profiler_trace.json
+
+# For MPI usage, each rank gets its own file
+ls -la ./profiler_results/pytorch_profiler_trace_rank*.json
 ```
 
 ## Environment Variables
@@ -28,5 +32,13 @@ The profiler uses PyTorch's modern `torch::profiler` API and automatically:
 - Creates the output directory if it doesn't exist
 - Profiles all forward pass operations in DeepPotPT and DeepSpinPT
 - Saves profiling results to a JSON file when the object is destroyed
+- Automatically includes MPI rank in filename when MPI is available and initialized
+
+## Output Files
+
+- **Single-rank or non-MPI usage**: `pytorch_profiler_trace.json`
+- **MPI usage**: `pytorch_profiler_trace_rank{rank}.json` (e.g., `pytorch_profiler_trace_rank0.json`, `pytorch_profiler_trace_rank1.json`)
+
+This ensures that each MPI rank saves its profiling data to a separate file, preventing conflicts in multi-rank simulations.
 
 This is intended for development and debugging purposes.

source/api_cc/include/common.h

@@ -172,6 +172,12 @@ void get_env_nthreads(int& num_intra_nthreads, int& num_inter_nthreads);
  **/
 void get_env_pytorch_profiler(bool& enable_profiler, std::string& output_dir);
 
+/**
+ * @brief Get MPI rank if MPI is available and initialized, otherwise return 0.
+ * @return The MPI rank or 0 if MPI is not available/initialized.
+ **/
+int get_mpi_rank();
+
 /**
  * @brief Dynamically load OP library. This should be called before loading
  * graphs.

source/api_cc/src/DeepPotPT.cc

@@ -142,8 +142,16 @@ DeepPotPT::~DeepPotPT() {
 #ifdef BUILD_PYTORCH
   if (profiler_enabled) {
     try {
-      // Save profiler results to file
-      std::string output_file = profiler_output_dir + "/pytorch_profiler_trace.json";
+      // Save profiler results to file with MPI rank if available
+      int rank = get_mpi_rank();
+      std::string output_file;
+      if (rank >= 0) {
+        // MPI is available and initialized, include rank in filename
+        output_file = profiler_output_dir + "/pytorch_profiler_trace_rank" + std::to_string(rank) + ".json";
+      } else {
+        // MPI not available or not initialized, use original filename
+        output_file = profiler_output_dir + "/pytorch_profiler_trace.json";
+      }
       profiler_result = torch::profiler::disableProfiler();
       if (profiler_result) {
         profiler_result->save(output_file);

source/api_cc/src/DeepSpinPT.cc

@@ -142,8 +142,16 @@ DeepSpinPT::~DeepSpinPT() {
 #ifdef BUILD_PYTORCH
   if (profiler_enabled) {
     try {
-      // Save profiler results to file
-      std::string output_file = profiler_output_dir + "/pytorch_profiler_trace.json";
+      // Save profiler results to file with MPI rank if available
+      int rank = get_mpi_rank();
+      std::string output_file;
+      if (rank >= 0) {
+        // MPI is available and initialized, include rank in filename
+        output_file = profiler_output_dir + "/pytorch_profiler_trace_rank" + std::to_string(rank) + ".json";
+      } else {
+        // MPI not available or not initialized, use original filename
+        output_file = profiler_output_dir + "/pytorch_profiler_trace.json";
+      }
       profiler_result = torch::profiler::disableProfiler();
       if (profiler_result) {
         profiler_result->save(output_file);

source/api_cc/src/common.cc

@@ -8,6 +8,13 @@
 #include <sstream>
 #include <string>
 
+// Try to include MPI if available - this will be a no-op if MPI is not available
+#ifdef __has_include
+  #if __has_include(<mpi.h>)
+    #include <mpi.h>
+  #endif
+#endif
+
 #include "AtomMap.h"
 #include "device.h"
 #if defined(_WIN32)
@@ -398,6 +405,20 @@ void deepmd::get_env_pytorch_profiler(bool& enable_profiler, std::string& output
   }
 }
 
+int deepmd::get_mpi_rank() {
+  int rank = -1;  // Use -1 to indicate MPI not available/initialized
+  // Try to detect MPI at runtime
+  #ifdef MPI_H
+  int initialized = 0;
+  if (MPI_Initialized(&initialized) == MPI_SUCCESS && initialized) {
+    if (MPI_Comm_rank(MPI_COMM_WORLD, &rank) != MPI_SUCCESS) {
+      rank = -1;  // fallback to -1 if MPI_Comm_rank fails
+    }
+  }
+  #endif
+  return rank;
+}
+
 static inline void _load_library_path(std::string dso_path) {
 #if defined(_WIN32)
   void* dso_handle = LoadLibrary(dso_path.c_str());

source/api_cc/tests/test_pytorch_profiler.cc

@@ -73,4 +73,10 @@ TEST_F(TestPyTorchProfiler, test_profiler_disabled_with_zero) {
 
   EXPECT_FALSE(enable_profiler);
   EXPECT_EQ(output_dir, "./profiler_output");
+}
+
+TEST_F(TestPyTorchProfiler, test_mpi_rank_detection) {
+  // Test that MPI rank detection returns valid rank (-1 when MPI not initialized, >= 0 when initialized)
+  int rank = deepmd::get_mpi_rank();
+  EXPECT_GE(rank, -1);  // Rank should be -1 (not available) or >= 0 (valid rank)
 }