improve compilation time and create timetrace plot (#773)

* improve compilation time and create time-trace plot

* set default value for nproc

[ROCm/rccl commit: f35777e9b0]

projects/rccl/CMakeLists.txt

@@ -273,6 +273,7 @@ set(SRC_FILES
   src/collectives/device/reduce_kernel.h
   src/collectives/device/reduce_scatter.h
   src/collectives/device/sendrecv.h
+  src/collectives/device/msccl_kernel_impl.h
   src/collectives/gather.cc
   src/collectives/msccl.cc
   src/collectives/reduce.cc
@@ -413,21 +414,22 @@ set(SRC_FILES
 if (BUILD_ALLREDUCE_ONLY)
   add_definitions(-DBUILD_ALLREDUCE_ONLY)
   set(CU_SOURCES
-      src/collectives/device/all_reduce.cu
+      # src/collectives/device/all_reduce.cu
       src/collectives/device/sendrecv.cu
       src/collectives/device/functions.cu
-      src/collectives/device/msccl_kernel.cu)
+      # src/collectives/device/msccl_kernel.cu
+      )
 else()
   set(CU_SOURCES
       src/collectives/device/all_gather.cu
-      src/collectives/device/all_reduce.cu
+      # src/collectives/device/all_reduce.cu
       src/collectives/device/alltoall_pivot.cu
       src/collectives/device/broadcast.cu
       src/collectives/device/functions.cu
-      src/collectives/device/msccl_kernel.cu
+      # src/collectives/device/msccl_kernel.cu
       src/collectives/device/onerank_reduce.cu
-      src/collectives/device/reduce.cu
-      src/collectives/device/reduce_scatter.cu
+      # src/collectives/device/reduce.cu
+      # src/collectives/device/reduce_scatter.cu
       src/collectives/device/sendrecv.cu)
 endif()
 list(APPEND SRC_FILES ${CU_SOURCES})
@@ -461,6 +463,11 @@ foreach(SRC_FILE ${SRC_FILES})
   )
 endforeach()
 
+expand_collectives("all_reduce" "AllReduce")
+expand_collectives("reduce" "Reduce")
+expand_collectives("reduce_scatter" "ReduceScatter")
+expand_collectives("msccl_kernel" "MscclKernel")
+
 # Create an initial git_version.cpp file (that will be updated with latest git version)
 #==================================================================================================
 file(WRITE ${CMAKE_CURRENT_BINARY_DIR}/git_version.cpp "")

projects/rccl/cmake/Dependencies.cmake

@@ -70,6 +70,54 @@ if(NOT GTest_FOUND AND BUILD_TESTS OR INSTALL_DEPENDENCIES)
     endif()
 endif()
 
+set(DATATYPES_INT
+"int8_t"
+"uint8_t"
+"int32_t"
+"uint32_t"
+"int64_t"
+"uint64_t"
+  )
+set(DATATYPES_FLOAT
+  "half"
+  "float"
+  "double"
+  "rccl_bfloat16"
+  )
+
+function(expand_collectives FILE FUNC)
+  set(REDOP Sum Prod Min Max PreMulSum SumPostDiv)
+  if (FUNC STREQUAL "MscclKernel")
+    set(REDOP_FILTERED Sum Prod Min Max)
+  else()
+    set(REDOP_FILTERED ${REDOP})
+  endif()
+  foreach(REDOP_CURRENT IN LISTS REDOP_FILTERED)
+    foreach(DATA_TYPE ${DATATYPES_INT} ${DATATYPES_FLOAT})
+      if (REDOP_CURRENT STREQUAL "SumPostDiv" AND DATA_TYPE IN_LIST DATATYPES_FLOAT)
+        continue()  # Skip the iteration for DATATYPES_FLOAT when REDOP_CURRENT is SumPostDiv
+      endif()
+      set(FILE_NAME "${HIPIFY_DIR}/src/collectives/device/${FILE}_${REDOP_CURRENT}_${DATA_TYPE}.cpp")
+      message(STATUS "Generating ${FILE_NAME}")
+      if (FUNC STREQUAL "MscclKernel")
+        file(WRITE ${FILE_NAME}
+          "#include \"${FILE}_impl.h\"
+          #include \"primitives.h\"
+          #include \"collectives.h\"
+          #include \"devcomm.h\"
+          MSCCL_IMPL_KERNEL_ENTRY_FUNC_DEVREDOP_TYPE(${REDOP_CURRENT}, ${DATA_TYPE});")
+      else()
+        file(WRITE ${FILE_NAME}
+          "#include \"${FILE}.h\"
+          #include \"common.h\"
+          #include \"collectives.h\"
+          IMPL_COLL3(${FUNC}, ${REDOP_CURRENT}, ${DATA_TYPE});")
+      endif()
+      list(APPEND HIP_SOURCES ${FILE_NAME})
+    endforeach()
+  endforeach()
+  set(HIP_SOURCES ${HIP_SOURCES} PARENT_SCOPE)
+endfunction()
 
 # Find or download/install rocm-cmake project
 set( PROJECT_EXTERN_DIR ${CMAKE_CURRENT_BINARY_DIR}/extern )

projects/rccl/install.sh

@@ -17,6 +17,7 @@ function display_help()
     echo "    -h|--help                  Prints this help message"
     echo "    -i|--install               Install RCCL library (see --prefix argument below)"
     echo "       --local_gpu_only        Only compile for local GPU architecture"
+    echo "       --max-jobs              Use nproc instead of default number of 16"
     echo "       --no_clean              Don't delete files if they already exist"
     echo "       --npkit-enable          Compile with npkit enabled"
     echo "    -p|--package_build         Build RCCL package"
@@ -26,6 +27,7 @@ function display_help()
     echo "    -r|--run_tests_quick       Run small subset of rccl unit tests (must be built already)"
     echo "       --static                Build RCCL as a static library instead of shared library"
     echo "    -t|--tests_build           Build rccl unit tests, but do not run"
+    echo "       --time-trace            Plot the build time of RCCL"
     echo "       --verbose               Show compile commands"
 }
 
@@ -50,6 +52,9 @@ run_tests_all=false
 build_static=false
 build_tests=false
 build_verbose=0
+time_trace=false
+enable_all_jobs=false
+enable_ninja=""
 
 # #################################################
 # Parameter parsing
@@ -58,7 +63,7 @@ build_verbose=0
 # check if we have a modern version of getopt that can handle whitespace and long parameters
 getopt -T
 if [[ $? -eq 4 ]]; then
-    GETOPT_PARSE=$(getopt --name "${0}" --longoptions address-sanitizer,build_allreduce_only,dependencies,debug,disable_backtrace,fast,help,install,local_gpu_only,no_clean,npkit-enable,package_build,prefix:,rm-legacy-include-dir,run_tests_all,run_tests_quick,tests_build,verbose --options hidptrs -- "$@")
+    GETOPT_PARSE=$(getopt --name "${0}" --longoptions address-sanitizer,build_allreduce_only,dependencies,debug,disable_backtrace,fast,help,install,local_gpu_only,no_clean,npkit-enable,package_build,prefix:,rm-legacy-include-dir,run_tests_all,run_tests_quick,tests_build,time-trace,max-jobs,verbose --options hidptrs -- "$@")
 else
     echo "Need a new version of getopt"
     exit 1
@@ -82,6 +87,7 @@ while true; do
     -h | --help)                     display_help;                               exit 0 ;;
     -i | --install)                  install_library=true;                       shift ;;
          --local_gpu_only)           build_local_gpu_only=true;                  shift ;;
+         --max-jobs)                 enable_all_jobs=true;                       shift ;;
          --no_clean)                 clean_build=false;                          shift ;;
          --npkit-enable)             npkit_enabled=true;                         shift ;;
     -p | --package_build)            build_package=true;                         shift ;;
@@ -91,6 +97,7 @@ while true; do
          --run_tests_all)            run_tests=true; run_tests_all=true;         shift ;;
          --static)                   build_static=true;                          shift ;;
     -t | --tests_build)              build_tests=true;                           shift ;;
+         --time-trace)               time_trace=true;                            shift ;;
          --verbose)                  build_verbose=1;                            shift ;;
     --) shift ; break ;;
     *)  echo "Unexpected command line parameter received; aborting";
@@ -291,17 +298,30 @@ fi
 
 check_exit_code "$?"
 
-if ($build_tests) || (($run_tests) && [[ ! -f ./test/rccl-UnitTests ]]); then
-    CXX=$ROCM_BIN_PATH/hipcc $cmake_executable $cmake_common_options -DBUILD_TESTS=ON -DNPKIT_FLAGS="${npkit_options}" -DCMAKE_INSTALL_PREFIX=$ROCM_PATH -DROCM_PATH=$ROCM_PATH ../../.
+if ($enable_all_jobs); then
+    job_number=$(nproc)
 else
-    CXX=$ROCM_BIN_PATH/hipcc $cmake_executable $cmake_common_options -DBUILD_TESTS=OFF -DNPKIT_FLAGS="${npkit_options}" -DCMAKE_INSTALL_PREFIX=$ROCM_PATH -DROCM_PATH=$ROCM_PATH ../../.
+    job_number=16
+fi
+
+if ($time_trace); then
+    build_system="ninja"
+    enable_ninja="-GNinja"
+else
+    build_system="make"
+fi
+
+if ($build_tests) || (($run_tests) && [[ ! -f ./test/rccl-UnitTests ]]); then
+    CXX=$ROCM_BIN_PATH/hipcc $cmake_executable $cmake_common_options -DBUILD_TESTS=ON -DNPKIT_FLAGS="${npkit_options}" -DCMAKE_INSTALL_PREFIX=$ROCM_PATH -DROCM_PATH=$ROCM_PATH $enable_ninja ../../.
+else
+    CXX=$ROCM_BIN_PATH/hipcc $cmake_executable $cmake_common_options -DBUILD_TESTS=OFF -DNPKIT_FLAGS="${npkit_options}" -DCMAKE_INSTALL_PREFIX=$ROCM_PATH -DROCM_PATH=$ROCM_PATH $enable_ninja ../../.
 fi
 check_exit_code "$?"
 
 if ($install_library); then
-    VERBOSE=${build_verbose} make -j$(nproc) install
+    VERBOSE=${build_verbose} $build_system -j $job_number install
 else
-    VERBOSE=${build_verbose} make -j$(nproc)
+    VERBOSE=${build_verbose} $build_system -j $job_number
 fi
 check_exit_code "$?"
 
@@ -323,3 +343,10 @@ if ($run_tests); then
         exit 1
     fi
 fi
+
+if ($time_trace) then
+    cd ../../tools/time-trace
+    chmod +x ./rccl-TimeTrace.sh
+    echo "Generating RCCL-compile-timeline.html..."
+    ./rccl-TimeTrace.sh
+fi

projects/rccl/src/collectives/device/all_reduce.cu

@@ -3,9 +3,10 @@
  *
  * See LICENSE.txt for license information
  ************************************************************************/
+/*This file is now generated in CMake*/
 
-#include "all_reduce.h"
-#include "common.h"
-#include "collectives.h"
+// #include "all_reduce.h"
+// #include "common.h"
+// #include "collectives.h"
 
-IMPL_COLL_R(AllReduce);
+// IMPL_COLL_R(AllReduce);

projects/rccl/src/collectives/device/common.h

@@ -594,37 +594,37 @@ __device__  __attribute__((noinline)) void NCCL_FUNC_NAME(func, algo, proto, dev
 #endif
 
 // Only generate inline kernels for LL
-#define IMPL_COLL4(func, algo, devredop, type, ncclType) \
+#define IMPL_COLL4(func, algo, devredop, type) \
   IMPL_COLL_FUNC(func, algo, LL,     devredop, type) \
   IMPL_COLL_FUNC(func, algo, LL128,  devredop, type) \
   IMPL_COLL_FUNC(func, algo, SIMPLE, devredop, type)
 
-#define IMPL_COLL3(func, devredop, type, ncclType) \
-  IMPL_COLL4(func, TREE,    devredop, type, ncclType) \
-  IMPL_COLL4(func, RING,    devredop, type, ncclType) \
-  IMPL_COLL4(func, COLLNET_DIRECT, devredop, type, ncclType) \
-  IMPL_COLL4(func, COLLNET_CHAIN, devredop, type, ncclType) \
-  IMPL_COLL4(func, NVLS, devredop, type, ncclType)
+#define IMPL_COLL3(func, devredop, type) \
+  IMPL_COLL4(func, TREE,    devredop, type) \
+  IMPL_COLL4(func, RING,    devredop, type) \
+  IMPL_COLL4(func, COLLNET_DIRECT, devredop, type) \
+  IMPL_COLL4(func, COLLNET_CHAIN, devredop, type) \
+  IMPL_COLL4(func, NVLS, devredop, type)
 
 #define IMPL_COLL2(func, devredop) \
-  IMPL_COLL3(func, devredop, int8_t,   ncclInt8) \
-  IMPL_COLL3(func, devredop, uint8_t,  ncclUint8) \
-  IMPL_COLL3(func, devredop, int32_t,  ncclInt32) \
-  IMPL_COLL3(func, devredop, uint32_t, ncclUint32) \
-  IMPL_COLL3(func, devredop, int64_t,  ncclInt64) \
-  IMPL_COLL3(func, devredop, uint64_t, ncclUint64) \
-  IMPL_COLL3(func, devredop, half,     ncclFloat16) \
-  IMPL_COLL3(func, devredop, float,    ncclFloat32) \
-  IMPL_COLL3(func, devredop, double,   ncclFloat64) \
-  IMPL_COLL3(func, devredop, rccl_bfloat16, ncclBfloat16)
+  IMPL_COLL3(func, devredop, int8_t) \
+  IMPL_COLL3(func, devredop, uint8_t) \
+  IMPL_COLL3(func, devredop, int32_t) \
+  IMPL_COLL3(func, devredop, uint32_t) \
+  IMPL_COLL3(func, devredop, int64_t) \
+  IMPL_COLL3(func, devredop, uint64_t) \
+  IMPL_COLL3(func, devredop, half) \
+  IMPL_COLL3(func, devredop, float) \
+  IMPL_COLL3(func, devredop, double) \
+  IMPL_COLL3(func, devredop, rccl_bfloat16)
 
 #define IMPL_COLL2A(func, devredop) \
-  IMPL_COLL3(func, devredop, int8_t,   ncclInt8) \
-  IMPL_COLL3(func, devredop, uint8_t,  ncclUint8) \
-  IMPL_COLL3(func, devredop, int32_t,  ncclInt32) \
-  IMPL_COLL3(func, devredop, uint32_t, ncclUint32) \
-  IMPL_COLL3(func, devredop, int64_t,  ncclInt64) \
-  IMPL_COLL3(func, devredop, uint64_t, ncclUint64)
+  IMPL_COLL3(func, devredop, int8_t) \
+  IMPL_COLL3(func, devredop, uint8_t) \
+  IMPL_COLL3(func, devredop, int32_t) \
+  IMPL_COLL3(func, devredop, uint32_t) \
+  IMPL_COLL3(func, devredop, int64_t) \
+  IMPL_COLL3(func, devredop, uint64_t)
 
 // Reduction define all functions
 #define IMPL_COLL_R(func) \
@@ -636,7 +636,7 @@ __device__  __attribute__((noinline)) void NCCL_FUNC_NAME(func, algo, proto, dev
   IMPL_COLL2A(func, SumPostDiv)
 
 // Copy primitives only define one function for copy
-#define IMPL_COLL_C(func) IMPL_COLL3(func, Sum, int8_t, ncclInt8);
+#define IMPL_COLL_C(func) IMPL_COLL3(func, Sum, int8_t);
 
 // Point-to-point primitives only have one function/kernel.
 #define IMPL_COLL_P(func) \

projects/rccl/src/collectives/device/msccl_kernel_impl.h

@@ -5,6 +5,8 @@
  *
  * See LICENSE.txt for license information
  ************************************************************************/
+#ifndef MSSCLKERNELIMPL_H
+#define MSSCLKERNELIMPL_H
 
 #include "devcomm.h"
 #include "primitives.h"
@@ -13,7 +15,7 @@
 #include "msccl/msccl_struct.h"
 #include "msccl/msccl_kernel.h"
 
-__shared__ struct mscclShmemData mscclShmem;
+extern __shared__ struct mscclShmemData mscclShmem;
 
 #define MSCCL_MAX_ITER 65536
 
@@ -386,4 +388,4 @@ __global__ void MSCCL_KERNEL_ENTRY_NAME(devredop, type, Simple)(struct ncclDevCo
   MSCCL_IMPL_KERNEL_ENTRY_FUNC_DEVREDOP(Max) \
   MSCCL_IMPL_KERNEL_ENTRY_FUNC_DEVREDOP(Min)
 
-MSCCL_IMPL_KERNEL_ENTRY_FUNC()
+#endif

projects/rccl/src/collectives/device/reduce.cu

@@ -4,8 +4,10 @@
  * See LICENSE.txt for license information
  ************************************************************************/
 
-#include "reduce.h"
-#include "common.h"
-#include "collectives.h"
+/*This file is now generated in CMake*/
 
-IMPL_COLL_R(Reduce);
+// #include "reduce.h"
+// #include "common.h"
+// #include "collectives.h"
+
+// IMPL_COLL_R(Reduce);

projects/rccl/src/collectives/device/reduce_scatter.cu

@@ -4,8 +4,10 @@
  * See LICENSE.txt for license information
  ************************************************************************/
 
-#include "reduce_scatter.h"
-#include "common.h"
-#include "collectives.h"
+/*This file is now generated in CMake*/
 
-IMPL_COLL_R(ReduceScatter);
+// #include "reduce_scatter.h"
+// #include "common.h"
+// #include "collectives.h"
+
+// IMPL_COLL_R(ReduceScatter);

projects/rccl/tools/time-trace/rccl-TimeTrace.sh

@@ -0,0 +1,50 @@
+#!/bin/bash
+
+# Directory path to search for JSON files
+directory="../../build/release"
+
+if command -v pip &>/dev/null; then
+    echo "pip is already installed."
+else
+    echo "pip is not installed. Installing..."
+    sudo apt-get update
+    sudo apt install python3-pip
+fi
+
+required_library='pandas'
+
+# Check if pandas is installed
+if python3 -c "import $required_library" &> /dev/null; then
+    echo "$required_library is already installed."
+else
+    echo "$required_library is not installed. Installing..."
+    pip3 install $required_library
+fi
+
+required_library='plotly'
+
+# Check if the library is installed
+if python3 -c "import $required_library" &> /dev/null; then
+    echo "$required_library is already installed."
+else
+    echo "$required_library is not installed. Installing..."
+    pip3 install $required_library
+fi
+
+# Check if the file exists
+if [ ! -f "$directory/.ninja_log" ]; then
+  echo "File '$directory/.ninja_log' does not exist."
+  exit 1
+fi
+
+declare -A unique_values
+
+# Use awk to compare and delete duplicates
+awk '!unique_values[$5]++' "$directory/.ninja_log" > temp_file.txt
+mv temp_file.txt "$directory/.ninja_log"
+
+# Rename the file with .csv extension
+mv "$directory/.ninja_log" "$directory/time_trace.log"
+
+# Run the python program
+python3 time_trace_generator.py --min_val 5 --include_linking

projects/rccl/tools/time-trace/time_trace_generator.py

@@ -0,0 +1,109 @@
+import os
+import random
+import pandas as pd
+import plotly.graph_objects as go
+import argparse
+
+# Specify the path to the .log file
+log_file = '../../build/release/time_trace.log'
+
+def generateRandomColors(df, colorList):
+
+    for _ in range(len(df)):
+        r = random.random()
+        g = random.random()
+        b = random.random()
+        colorList.append(f"rgb({int(r*255)}, {int(g*255)}, {int(b*255)})")
+    
+    return colorList
+
+def plotCompileTime(log_file, minVal):
+    colors = []
+
+    # read the log file and extract the data from it
+    # st:   start-time (ms)
+    # et:   end-time (ms)
+    # ts:   timestamp
+    # file: path to file
+    # hash: command hash 
+    df = pd.read_csv(log_file, delimiter='\t', header=None, 
+                     names=['st', 'et', 'ts', 'file', 'hash'])
+    df = df.iloc[1:]
+
+    # include file name only
+    df['file'] = df['file'].apply(os.path.basename)
+
+    # convert to seconds
+    df['st'] = df['st'].astype(int) / 1000
+    df['et'] = df['et'].astype(int) / 1000
+
+    # calculate compilation duration of the file
+    df['dur'] = df['et'].astype(int) - df['st'].astype(int)
+
+    if args.include_linking == 0:
+        # drop the last two rows which are related to linking
+        df = df.drop(df.index[-2:]) 
+    
+    # if minVal specified remove the rows from the df where df['dur'] < minVal
+    df = df[df['dur'] >= minVal]
+
+    maxEt = int(df['et'].max())
+    df = df[::-1]   # reverse df
+
+    colors = generateRandomColors(df, colors)
+
+    fig = go.Figure(go.Bar(
+        y=df['file'],
+        x=df['dur'],
+        orientation='h',
+        marker=dict(color=colors),
+        base=df['st'],
+        textposition='auto',
+        customdata=df['dur'],
+        hovertemplate='<b>Time:</b> %{customdata} seconds<br>' +
+                    '<b>File Name:</b> %{y}<br>'
+    ))
+
+    # Customize the layout
+    fig.update_layout(
+        title="RCCL TOTAL COMPILE TIME LINE",
+        xaxis_title='Duration (seconds)',
+        yaxis_title='file name',
+        bargap=0.1,
+        plot_bgcolor='#36454F',  # Set the plot background color to black
+        paper_bgcolor='#36454F',  # Set the paper background color to black
+        font=dict(
+            family="Arial",
+            size=11,
+            color="white"
+        ),
+    )
+
+    # add custom text annotation at the top right corner
+    fig.update_layout(
+        annotations=[
+            go.layout.Annotation(
+                x=1,
+                y=1,
+                xref="paper",
+                yref="paper",
+                text="Total Time: "+ str(maxEt) + " seconds",
+                showarrow=False,
+                font=dict(
+                    size=18,
+                    color="white"
+                )
+            )
+        ]
+    )
+
+    # convert the plot to an html file
+    fig.write_html("RCCL-compile-timeline.html", auto_open=False)
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--min_val", nargs='?', default='5', type=int, help="Ignore any if it's less than the value provided.")
+    parser.add_argument("--include_linking", action='store_true', help="Include linking when plotting.")
+    args = parser.parse_args()
+    
+    plotCompileTime(log_file, args.min_val)