Populate in/out split size info for all_to_all from CPU to CUDA kernel (

#822)

Summary:
X-link: pytorch/pytorch#112308

Pull Request resolved: #822

This diff populates all_to_all input and out split size from CPU op to GPU kernel when valid.

bypass-github-pytorch-ci-checks

Reviewed By: aaronenyeshi, idning

Differential Revision: D50762093

fbshipit-source-id: a118b9e2623ca0ac6b5f9e30cd554666a4c01a12

libkineto/src/output_json.cpp

@@ -36,6 +36,8 @@ static constexpr const char* kDtype = "dtype";
 static constexpr const char* kInMsgSize = "In msg size";
 static constexpr const char* kOutMsgSize = "Out msg size";
 static constexpr const char* kGroupSize = "Group size";
+static constexpr const char* kInSplit = "In split size";
+static constexpr const char* kOutSplit = "Out split size";
 
 #ifdef __linux__
 static constexpr char kDefaultLogFileFmt[] =
@@ -329,6 +331,20 @@ void ChromeTraceLogger::handleActivity(
           kDtype,
           dtype));
     }
+    // In/out split size are valid for all_to_all
+    const auto& inSplitSize = collectiveRecord->getMetadataValue(kInSplit);
+    const auto& outSplitSize = collectiveRecord->getMetadataValue(kOutSplit);
+    if (!inSplitSize.empty() && !outSplitSize.empty()) {
+      if (!arg_values.empty()) {
+        arg_values.append(",");
+      }
+      arg_values.append(fmt::format(
+          "\"{}\": {}, \"{}\": {}",
+          kInSplit,
+          inSplitSize,
+          kOutSplit,
+          outSplitSize));
+    }
   }
 
   std::string args = "";