From 59506b7fdbda0a4223dfe0574789e59e4b511f21 Mon Sep 17 00:00:00 2001
From: Yue Dong <yoyoyod@meta.com>
Date: Tue, 7 Nov 2023 01:40:38 -0800
Subject: [PATCH] Populate in/out split size info for all_to_all from CPU to
 CUDA kernel (#822)

Summary:
X-link: https://github.com/pytorch/pytorch/pull/112308

Pull Request resolved: https://github.com/pytorch/kineto/pull/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 | 16 ++++++++++++++++
 1 file changed, 16 insertions(+)

diff --git a/libkineto/src/output_json.cpp b/libkineto/src/output_json.cpp
index 339dde9e..78b3e966 100644
--- a/libkineto/src/output_json.cpp
+++ b/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 = "";