Add stack size UT (#1081)

* Add stack size UT

[ROCm/rccl commit: dc2d486ba0]

projects/rccl/test/CMakeLists.txt

@@ -68,9 +68,19 @@ if(BUILD_TESTS)
   endif()
 
   add_executable(rccl-UnitTests ${COMMON_SOURCE_FILES} ${TEST_SOURCE_FILES})
+
+  ## Set rccl-UnitTests include directories
   target_include_directories(rccl-UnitTests PRIVATE ${ROCM_PATH} ${GTEST_INCLUDE_DIRS})
-  target_include_directories(rccl-UnitTests PRIVATE ${PROJECT_BINARY_DIR}/include)        # for generated rccl.h header
+  target_include_directories(rccl-UnitTests PRIVATE ${PROJECT_BINARY_DIR}/include)            # for generated rccl.h header
   target_include_directories(rccl-UnitTests PRIVATE ${PROJECT_BINARY_DIR}/hipify/src/include) # for rccl_bfloat16.h
+
+  ## Set rccl-UnitTests compile definitions
+  if(LL128_ENABLED)
+    target_compile_definitions(rccl-UnitTests PRIVATE ENABLE_LL128)
+  endif()
+  target_compile_definitions(rccl-UnitTests PRIVATE ROCM_PATH="${ROCM_PATH}")
+
+  ## Set rccl-UnitTests linked libraries
   target_link_libraries(rccl-UnitTests PRIVATE ${GTEST_BOTH_LIBRARIES})
   target_link_libraries(rccl-UnitTests PRIVATE hip::host hip::device hsa-runtime64::hsa-runtime64)
   target_link_libraries(rccl-UnitTests PRIVATE Threads::Threads)

projects/rccl/test/StandaloneTests.cpp

@@ -9,7 +9,11 @@
 
 #include "StandaloneUtils.hpp"
 
-namespace RcclUnitTesting {
+namespace RcclUnitTesting 
+{
+  /**
+   * \brief Verify that each device is assigned to the right rank using ncclCommSplit API.
+   * ******************************************************************************************/
   TEST(Standalone, SplitComms_RankCheck)
   {
     // Check for multi-gpu
@@ -52,6 +56,9 @@ namespace RcclUnitTesting {
       NCCLCHECK(ncclCommDestroy(comm));
   }
 
+  /**
+   * \brief Creates a communicator for each device and gathers them all in one rank.
+   * ******************************************************************************************/
   TEST(Standalone, SplitComms_OneColor)
   {
     // Check for multi-gpu
@@ -93,6 +100,9 @@ namespace RcclUnitTesting {
       NCCLCHECK(ncclCommDestroy(comm));
   }
 
+  /**
+   * \brief Creates a communicator for each device and reduces them into (numDevices / 2) ranks.
+   * ******************************************************************************************/
   TEST(Standalone, SplitComms_Reduce)
   {
     // Check for multi-gpu
@@ -140,7 +150,10 @@ namespace RcclUnitTesting {
     for (auto& comm : comms)
       NCCLCHECK(ncclCommDestroy(comm));
   }
-
+  
+  /**
+   * \brief Verify there is no regression in timing for each protocol [LL, LL128, Simple]
+   * ******************************************************************************************/
   TEST(Standalone, RegressionTiming)
   {
     // timing
@@ -241,4 +254,41 @@ namespace RcclUnitTesting {
     else
       unsetenv("NCCL_PROTO");
   }
+
+  /**
+   * \brief Verify rccl generic kernel stack size for each gfx architecture is less than the
+   * expected MAX_STACK_SIZE.
+   * ******************************************************************************************/
+  TEST(Standalone, StackSize) {
+    const char* mainKernel = "rccl_main_kernel";
+
+    // Look for the .co files
+    std::vector<std::string> coFileList = splitString(executeCommand("find ../ -type f -name \"*.co\""), '\n');
+
+    // Check if the .co files exist in the build directory
+    if (coFileList.empty())
+      GTEST_SKIP() << "Skipping... Could not found required files in the build directory.";
+
+    for (const auto& file : coFileList) {
+      // Store the output in a list
+      std::string cmd = std::string(ROCM_PATH) + "/llvm/bin/llvm-readelf --notes " + file;
+      std::vector<std::string> metadata = splitString(executeCommand(cmd.c_str()), '\n');
+
+      // Skip if llvm is not installed
+      if (metadata.empty())
+        GTEST_SKIP() << "Skipping... llvm is not found.";
+
+      // Parse metadata from file and store it for each arch
+      ArchInfo archInfo = parseMetadata(metadata);
+
+      // iterate over each archs kernels
+      for (const auto& kernel : archInfo.kernels) {
+        if (kernel.name.find(mainKernel) != std::string::npos) {
+          // Kernel stack size should be less than or equal to the maxStackSize value
+          printf("[ INFO     ] Arch: %s Kernel: %s Size: %d\n", archInfo.archName.c_str(), kernel.name.c_str(), kernel.privateSegmentFixedSize);
+          EXPECT_LE(kernel.privateSegmentFixedSize, archInfo.archName == "gfx90a" ? MAX_STACK_SIZE_gfx90a : MAX_STACK_SIZE);
+        }
+      }
+    }
+  }
 }

projects/rccl/test/common/StandaloneUtils.hpp

@@ -1,6 +1,10 @@
 #ifndef STANDALONE_UTILS_H
 #define STANDALONE_UTILS_H
 
+#include <iostream>
+#include <cstdio>
+#include <regex>
+
 #define HIPCALL(cmd)                                                                          \
     do {                                                                                      \
         hipError_t error = (cmd);                                                             \
@@ -20,4 +24,83 @@
     }                                                           \
 } while(0)
 
+#define MAX_STACK_SIZE 112
+
+#ifdef ENABLE_LL128
+#define MAX_STACK_SIZE_gfx90a 288
+#else
+#define MAX_STACK_SIZE_gfx90a MAX_STACK_SIZE
+#endif
+
+struct KernelInfo {
+    std::string name;
+    int privateSegmentFixedSize = 0;
+};
+
+struct ArchInfo {
+    std::string archName;
+    std::vector<KernelInfo> kernels;
+};
+
+std::string executeCommand(const char* cmd) {
+    std::string result;
+    FILE* pipe = popen(cmd, "r");
+
+    if (!pipe) {
+        std::cerr << "Error executing command: " << cmd << std::endl;
+        return result;
+    }
+
+    char buffer[128];
+    while (!feof(pipe)) {
+        if (fgets(buffer, 128, pipe) != NULL) {
+            result += buffer;
+        }
+    }
+
+    pclose(pipe);
+    return result;
+}
+
+std::vector<std::string> splitString(const std::string& str, char delimiter) {
+    std::vector<std::string> result;
+    std::istringstream iss(str);
+
+    std::string line;
+    while(std::getline(iss, line, delimiter)) {
+        result.push_back(line);
+    }
+
+    return result;
+}
+
+
+ArchInfo parseMetadata(const std::vector<std::string>& list) {
+    ArchInfo archInfo;
+    KernelInfo currKernelInfo;
+    
+    std::regex amdhsaTargetRegex("amdhsa.target:\\s+(?:'?)amdgcn-amd-amdhsa--(\\w+)(?:'?)");
+    std::regex kernelNameRegex("\\.name:\\s+(\\w+)");
+    std::regex privateSegmentSizeRegex("\\.private_segment_fixed_size:\\s+(\\d+)");
+    
+    for (const auto& line : list) {
+        std::smatch match;
+
+        if (std::regex_search(line, match, amdhsaTargetRegex)) {
+            archInfo.archName = match[1];
+        } else if (std::regex_search(line, match, kernelNameRegex)) {
+            currKernelInfo.name = match[1];
+        } else if (std::regex_search(line, match, privateSegmentSizeRegex)) {
+            currKernelInfo.privateSegmentFixedSize = std::stoi(match[1]);
+        }
+        
+        if (!currKernelInfo.name.empty() && currKernelInfo.privateSegmentFixedSize != 0) {
+            archInfo.kernels.push_back(currKernelInfo);
+            currKernelInfo = {}; // Empty kernelInfo
+        }
+    }
+    
+    return archInfo;
+}
+
 #endif