Undo revert of the PR #1014 (#1017)

Summary:
Pull Request resolved: #1017

Reverts a rollback D66458621

Revert the Kineto rollback, this would have partially solved the issue since this part controls transmission of the metadata to the corresponding kernel. but record_param_comms in pytorch is the real issue and was still recording this metadata and would still make an invalid trace JSON when working with GPUs>30 (our truncation case).

Reviewed By: sraikund16

Differential Revision: D66475394

fbshipit-source-id: 5b781ccb27fa898a1a6496c72733f72fd31c822e

libkineto/src/output_json.cpp

@@ -38,6 +38,8 @@ static constexpr const char* kOutSplit = "Out split size";
 static constexpr const char* kProcessGroupName = "Process Group Name";
 static constexpr const char* kProcessGroupDesc = "Process Group Description";
 static constexpr const char* kGroupRanks = "Process Group Ranks";
+static constexpr const char* kInTensorsStart = "Input Tensors start";
+static constexpr const char* kOutTensorsStart = "Output Tensors start";
 static constexpr const char* kRank = "Rank";
 static constexpr const char* kP2pSrc = "Src Rank";
 static constexpr const char* kP2pDst = "Dst Rank";
@@ -419,6 +421,24 @@ void ChromeTraceLogger::handleActivity(const libkineto::ITraceActivity& op) {
           kDtype,
           dtype));
     }
+    const auto& input_tensor_starts =
+        collectiveRecord->getMetadataValue(kInTensorsStart);
+    const auto output_tensor_starts =
+        collectiveRecord->getMetadataValue(kOutTensorsStart);
+    if (!input_tensor_starts.empty()) {
+      if (!arg_values.empty()) {
+        arg_values.append(",");
+      }
+      arg_values.append(
+          fmt::format(" \"{}\": {}", kInTensorsStart, input_tensor_starts));
+    }
+    if (!output_tensor_starts.empty()) {
+      if (!arg_values.empty()) {
+        arg_values.append(",");
+      }
+      arg_values.append(
+          fmt::format(" \"{}\": {}", kOutTensorsStart, output_tensor_starts));
+    }
     // In/out split size are valid for all_to_all
     const auto& inSplitSize = collectiveRecord->getMetadataValue(kInSplit);
     const auto& outSplitSize = collectiveRecord->getMetadataValue(kOutSplit);