/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/

#include "tensorflow/core/profiler/convert/xplane_to_tf_data_stats.h"

#include "absl/container/flat_hash_map.h"
#include "absl/container/flat_hash_set.h"
#include "absl/strings/str_format.h"
#include "absl/strings/str_split.h"
#include "absl/strings/string_view.h"
#include "tensorflow/core/lib/gtl/map_util.h"
#include "tensorflow/core/platform/protobuf.h"
#include "tensorflow/core/profiler/protobuf/tf_data_stats.pb.h"
#include "tensorflow/core/profiler/utils/group_events.h"
#include "tensorflow/core/profiler/utils/html_utils.h"
#include "tensorflow/core/profiler/utils/tf_op_utils.h"
#include "tensorflow/core/profiler/utils/tf_xplane_visitor.h"
#include "tensorflow/core/profiler/utils/timespan.h"
#include "tensorflow/core/profiler/utils/xplane_schema.h"
#include "tensorflow/core/profiler/utils/xplane_visitor.h"

namespace tensorflow {
namespace profiler {

// 50 us from https://www.tensorflow.org/guide/data_performance_analysis
const int64 kSlowCallThresholdPs = 50 * 1000000;

namespace {

// Returns true if the given iterator event is for a root iterator.
bool IsRootIteratorEvent(const XEventVisitor& iterator_event) {
  std::vector<absl::string_view> split_result =
      absl::StrSplit(iterator_event.Name(), "::");
  // The root iterator's name contains only its own name (no parent
  // information).
  return split_result.size() == 2;
}

// Returns true if the given iterator event name is for an async iterator.
bool IsAsyncIterator(absl::string_view iterator_event_name) {
  static auto* kAsyncIterators = new absl::flat_hash_set<absl::string_view>(
      {"Prefetch", "ParallelInterleave", "ParallelMap", "ParseExample",
       "MapAndBatch", "DataService", "LegacyParallelInterleave"});
  return kAsyncIterators->contains(iterator_event_name);
}

void SetIteratorMetadata(int64 id, const XEventVisitor& event,
                         IteratorMetadata* metadata) {
  metadata->set_id(id);
  auto parent_id_stat = event.GetStat(StatType::kParentId);
  if (parent_id_stat.has_value()) {
    metadata->set_parent_id(parent_id_stat->IntValue());
  }
  metadata->set_name(IteratorName(event.Name()));
  metadata->set_long_name(event.Name().data(), event.Name().size());
  metadata->set_is_async(IsAsyncIterator(metadata->name()));
  // TODO(b/161831651): Set params.
}

// Returns the parent iterator's id if it is a root of a device input
// pipeline.
absl::optional<int64> FindDeviceInputPipeline(const XEventVisitor& event) {
  if (event.Type() == HostEventType::kDeviceInputPipelineSecondIterator) {
    auto parent_id_stat = event.GetStat(StatType::kParentId);
    if (parent_id_stat.has_value()) return parent_id_stat->IntValue();
  }
  return absl::nullopt;
}

// Processes EventForest to do the following:
// (1) set iterator metadata
// (2) find root iterator events
// (3) find device input pipeline ids
void ProcessEventForest(const EventForest& event_forest,
                        absl::flat_hash_set<int64>* device_input_pipeline_ids,
                        absl::flat_hash_map<int64, std::vector<EventNode*>>*
                            root_iterator_event_map,
                        TfDataStats* tf_data_stats) {
  const EventNodeMap& event_node_map = event_forest.GetEventNodeMap();
  auto iterator_event_list =
      gtl::FindOrNull(event_node_map, HostEventType::kIterator);
  if (!iterator_event_list) return;
  for (const auto& iterator_event : *iterator_event_list) {
    const XEventVisitor& iterator_event_visitor =
        iterator_event->GetEventVisitor();
    auto iterator_id_stat = iterator_event_visitor.GetStat(StatType::kStepId);
    if (!iterator_id_stat.has_value()) continue;
    int64 iterator_id = iterator_id_stat->IntValue();
    auto result = tf_data_stats->mutable_iterator_metadata()->insert(
        {iterator_id, IteratorMetadata()});
    IteratorMetadata& metadata = result.first->second;
    if (result.second) {
      // First time processing this iterator.
      SetIteratorMetadata(iterator_id, iterator_event_visitor, &metadata);
    }
    if (IsRootIteratorEvent(iterator_event_visitor)) {
      // Record root iterator events.
      (*root_iterator_event_map)[iterator_id].push_back(iterator_event.get());
    }
  }
  auto device_input_pipeline_second_iterator_events = gtl::FindOrNull(
      event_node_map, HostEventType::kDeviceInputPipelineSecondIterator);
  if (!device_input_pipeline_second_iterator_events) return;
  for (const auto& iterator_event :
       *device_input_pipeline_second_iterator_events) {
    const XEventVisitor& iterator_event_visitor =
        iterator_event->GetEventVisitor();
    auto iterator_id_stat = iterator_event_visitor.GetStat(StatType::kStepId);
    if (!iterator_id_stat.has_value()) continue;
    int64 iterator_id = iterator_id_stat->IntValue();
    auto result = tf_data_stats->mutable_iterator_metadata()->insert(
        {iterator_id, IteratorMetadata()});
    IteratorMetadata& metadata = result.first->second;
    if (result.second) {
      // First time processing this iterator.
      SetIteratorMetadata(iterator_id, iterator_event_visitor, &metadata);
      // Find and record device input pipeline ids.
      absl::optional<int64> device_input_pipeline_id =
          FindDeviceInputPipeline(iterator_event_visitor);
      if (device_input_pipeline_id.has_value()) {
        device_input_pipeline_ids->insert(*device_input_pipeline_id);
      }
    }
  }
}

void SetInputPipelineMetadata(int64 id, int64 name_id,
                              bool is_device_input_pipeline,
                              InputPipelineMetadata* metadata) {
  constexpr absl::string_view kHostInputPipelinePrefix = "Host:";
  constexpr absl::string_view kDeviceInputPipelinePrefix = "Device:";
  metadata->set_id(id);
  if (is_device_input_pipeline) {
    metadata->set_type(InputPipelineMetadata::DEVICE);
    metadata->set_name(absl::StrCat(kDeviceInputPipelinePrefix, name_id));
  } else {
    metadata->set_type(InputPipelineMetadata::HOST);
    metadata->set_name(absl::StrCat(kHostInputPipelinePrefix, name_id));
  }
}

void ProcessIteratorEvent(const EventNode& iterator_event,
                          InputPipelineStat* input_pipeline_stat,
                          bool is_blocking) {
  const XEventVisitor& visitor = iterator_event.GetEventVisitor();
  auto iterator_id_stat = visitor.GetStat(StatType::kStepId);
  if (!iterator_id_stat.has_value()) return;
  int64 iterator_id = iterator_id_stat->IntValue();
  auto result = input_pipeline_stat->mutable_iterator_stats()->insert(
      {iterator_id, IteratorStat()});
  IteratorStat& iterator_stat = result.first->second;
  if (result.second) {
    iterator_stat.set_id(iterator_id);
    iterator_stat.set_start_time_ps(visitor.TimestampPs());
  }
  iterator_stat.set_duration_ps(iterator_stat.duration_ps() +
                                visitor.DurationPs());
  int64 self_time_ps = visitor.DurationPs();
  tensorflow::profiler::Timespan self_time_span = visitor.GetTimespan();
  for (EventNode* child : iterator_event.GetChildren()) {
    const XEventVisitor& child_visitor = child->GetEventVisitor();
    if (ParseTfOpFullname(child_visitor.Name()).category == Category::kTfData) {
      int64 overlap_duration_ps =
          self_time_span.OverlappedDurationPs(child_visitor.GetTimespan());
      ProcessIteratorEvent(*child, input_pipeline_stat,
                           is_blocking && overlap_duration_ps);
      // Note: Assume no overlap between child events.
      self_time_ps -= overlap_duration_ps;
    }
  }
  iterator_stat.set_self_time_ps(iterator_stat.self_time_ps() + self_time_ps);
  iterator_stat.set_is_blocking(iterator_stat.is_blocking() || is_blocking);
  iterator_stat.set_num_calls(iterator_stat.num_calls() + 1);
}

void SetBottleneckIteratorId(InputPipelineStat* input_pipeline_stat) {
  int64 bottleneck_iterator_id = 0;
  int64 max_self_time = 0;
  for (const auto& pair : input_pipeline_stat->iterator_stats()) {
    const auto& id = pair.first;
    const auto& iterator_stat = pair.second;
    if (iterator_stat.is_blocking() &&
        iterator_stat.self_time_ps() > max_self_time) {
      bottleneck_iterator_id = id;
      max_self_time = iterator_stat.self_time_ps();
    }
  }
  input_pipeline_stat->set_bottleneck_iterator_id(bottleneck_iterator_id);
}

void ProcessInputPipelines(
    const absl::flat_hash_set<int64>& device_input_pipeline_ids,
    absl::flat_hash_map<int64, std::vector<EventNode*>>*
        root_iterator_event_map,
    TfDataStats* tf_data_stats) {
  auto* input_pipelines = tf_data_stats->mutable_input_pipelines();
  int64 num_host_input_pipelines = 0;
  int64 num_device_input_pipelines = 0;
  for (auto& id_and_events : *root_iterator_event_map) {
    auto& root_iterator_id = id_and_events.first;
    auto& root_iterator_events = id_and_events.second;
    absl::c_sort(root_iterator_events,
                 [](const EventNode* lhs, const EventNode* rhs) {
                   return lhs->GetEventVisitor().DurationPs() >
                          rhs->GetEventVisitor().DurationPs();
                 });
    auto result =
        input_pipelines->insert({root_iterator_id, InputPipelineStats()});
    InputPipelineStats& input_pipeline_stats = result.first->second;
    InputPipelineMetadata* metadata = input_pipeline_stats.mutable_metadata();
    if (result.second) {
      bool is_device_input_pipeline =
          device_input_pipeline_ids.contains(root_iterator_id);
      int64 name_id = is_device_input_pipeline ? num_device_input_pipelines++
                                               : num_host_input_pipelines++;
      SetInputPipelineMetadata(root_iterator_id, name_id,
                               is_device_input_pipeline, metadata);
    }
    int64 sum_latency_ps = 0;
    int64 min_latency_ps = INT64_MAX;
    int64 max_latency_ps = 0;
    int64 num_slow_calls = 0;
    for (const EventNode* root_iterator_event : root_iterator_events) {
      InputPipelineStat* stat = input_pipeline_stats.add_stats();
      ProcessIteratorEvent(*root_iterator_event, stat,
                           /*is_blocking*/ true);
      SetBottleneckIteratorId(stat);
      int64 latency_ps = root_iterator_event->GetEventVisitor().DurationPs();
      sum_latency_ps += latency_ps;
      min_latency_ps = std::min(min_latency_ps, latency_ps);
      max_latency_ps = std::max(max_latency_ps, latency_ps);
      if (latency_ps > kSlowCallThresholdPs) num_slow_calls++;
    }
    input_pipeline_stats.set_avg_latency_ps(sum_latency_ps /
                                            root_iterator_events.size());
    input_pipeline_stats.set_min_latency_ps(min_latency_ps);
    input_pipeline_stats.set_max_latency_ps(max_latency_ps);
    input_pipeline_stats.set_num_slow_calls(num_slow_calls);
  }
}

void SetBottleneckAnalysis(absl::string_view host_name,
                           const TfDataStats& tf_data_stats,
                           TfDataBottleneckAnalysis* bottleneck_analysis) {
  for (const auto& id_and_stats : tf_data_stats.input_pipelines()) {
    const InputPipelineStats& input_pipeline_stats = id_and_stats.second;
    if (input_pipeline_stats.metadata().type() ==
            InputPipelineMetadata::DEVICE ||
        input_pipeline_stats.max_latency_ps() <=
            bottleneck_analysis->max_latency_ps()) {
      // Ignore device input pipelines and input pipelines faster than the
      // current bottleneck.
      continue;
    }
    bottleneck_analysis->set_host(host_name.data(), host_name.size());
    bottleneck_analysis->set_input_pipeline(
        input_pipeline_stats.metadata().name());
    bottleneck_analysis->set_max_latency_ps(
        input_pipeline_stats.max_latency_ps());
    const IteratorMetadata& metadata = tf_data_stats.iterator_metadata().at(
        input_pipeline_stats.stats(0).bottleneck_iterator_id());
    bottleneck_analysis->set_iterator_name(metadata.name());
    bottleneck_analysis->set_iterator_long_name(metadata.long_name());
  }
}

std::string GetSuggestion(BottleneckType type) {
  constexpr absl::string_view kPlaybookLink =
      "https://www.tensorflow.org/guide/data_performance_analysis";
  constexpr absl::string_view kPlaybookSourceDatasetLink =
      "https://www.tensorflow.org/guide/"
      "data_performance_analysis#source_datasets";
  constexpr absl::string_view kPlaybookCpuUtilizationLink =
      "https://www.tensorflow.org/guide/"
      "data_performance_analysis#3_are_you_reaching_high_cpu_utilization";
  constexpr absl::string_view kPlaybookTransformationLink =
      "https://www.tensorflow.org/guide/"
      "data_performance_analysis#transformation_datasets";
  constexpr absl::string_view kTfGuideParallelDataExtractionLink =
      "https://www.tensorflow.org/guide/"
      "data_performance#parallelizing_data_extraction";
  constexpr absl::string_view kTfGuideParallelTransformationLink =
      "https://www.tensorflow.org/guide/"
      "data_performance#parallelizing_data_transformation";
  constexpr absl::string_view kTfGuideCacheLink =
      "https://www.tensorflow.org/guide/data_performance#caching";
  switch (type) {
    case BottleneckType::kSlowSource:
      return absl::StrFormat(
          "1. Check the locality of a host and input data. Ideally, they "
          "should be in the same cell (or very close, like the same "
          "region).<br/>"
          "2. Parallelize reading from this dataset source. See %s and %s for "
          "more details.<br/>",
          AnchorElement(kPlaybookSourceDatasetLink, "here"),
          AnchorElement(kTfGuideParallelDataExtractionLink, "here"));
    case BottleneckType::kSlowRemoteSource:
      return absl::StrFormat(
          "1. The remote data source is slow. Profile its host to analyze the "
          "issue further.<br/>"
          "2. See %s for other suggestions.",
          AnchorElement(kPlaybookLink, "this"));
    case BottleneckType::kSlowTransformationWithParallelVersion:
      return absl::StrFormat(
          "1. Parallelize this transformation by setting "
          "<code>num_parallel_calls=tf.data.experimental.AUTOTUNE</code>. See "
          "%s for more details.<br/>"
          "2. Consider adding <code>cache</code> after this transformation if "
          "your data fits into memory and it is appropriate (e.g., there is no "
          "randomness in upstream transformations like <code>shuffle</code>). "
          "See %s for more details.<br/>"
          "3. Find more resources %s.",
          AnchorElement(kTfGuideParallelTransformationLink, "this"),
          AnchorElement(kTfGuideCacheLink, "this"),
          AnchorElement(kPlaybookTransformationLink, "here"));
    case BottleneckType::kSlowTransformationWithoutParallelVersion:
      return absl::StrFormat(
          "1. This transformation is inherently sequential. Add outer "
          "parallelism by running multiple copies of the input pipeline over "
          "sharded inputs and combining the results. See %s for more "
          "details.<br/>"
          "2. Consider adding <code>cache</code> after this transformation if "
          "your data fits into memory and it is appropriate (e.g., there is no "
          "randomness in upstream transformations like <code>shuffle</code>). "
          "See %s for more details.<br/>"
          "3. Find more resources %s.",
          AnchorElement(kPlaybookTransformationLink, "this"),
          AnchorElement(kTfGuideCacheLink, "this"),
          AnchorElement(kPlaybookCpuUtilizationLink, "here"));
    default:
      return absl::StrFormat("See %s for suggestions.",
                             AnchorElement(kPlaybookLink, "this"));
  }
}

void SetSuggestion(TfDataBottleneckAnalysis* bottleneck_analysis) {
  if (bottleneck_analysis->max_latency_ps() <= kSlowCallThresholdPs) return;
  bottleneck_analysis->set_suggestion(
      GetSuggestion(GetBottleneckType(bottleneck_analysis->iterator_name())));
}

void SetSummary(CombinedTfDataStats* combined_tf_data_stats) {
  int64 max_latency_ps =
      combined_tf_data_stats->bottleneck_analysis().max_latency_ps();
  if (max_latency_ps > kSlowCallThresholdPs) {
    combined_tf_data_stats->set_is_input_bound(true);
    combined_tf_data_stats->set_summary(
        "Your profile has a tf.data input pipeline slower than 50 us. Below "
        "shows a bottleneck in the slow input pipeline and a suggestion on how "
        "to fix it.");
  } else if (max_latency_ps > 0) {
    combined_tf_data_stats->set_is_input_bound(false);
    combined_tf_data_stats->set_summary(
        "Your profile does not have any tf.data input pipeline slower than 50 "
        "us. Your job could be still input bound if this profile didn't "
        "capture all workers.");
  } else {
    combined_tf_data_stats->set_is_input_bound(false);
    combined_tf_data_stats->set_summary(
        "No tf.data activitiy captured in your profile. If your job uses "
        "tf.data, try to capture a longer profile.");
  }
}

}  // namespace

BottleneckType GetBottleneckType(absl::string_view bottleneck_iterator_name) {
  static auto* kBottleneckTypeMap = new absl::flat_hash_map<absl::string_view,
                                                            BottleneckType>(
      {// Read from storage.
       {"TFRecord", BottleneckType::kSlowSource},
       {"SSTable", BottleneckType::kSlowSource},
       {"RecordIO", BottleneckType::kSlowSource},
       {"Spanner", BottleneckType::kSlowSource},
       {"TFColumn", BottleneckType::kSlowSource},
       {"SleepwalkRemoteDataset", BottleneckType::kSlowSource},
       {"TextLine", BottleneckType::kSlowSource},
       {"StitchedTimelineDataset", BottleneckType::kSlowSource},
       {"DateKeyDataset", BottleneckType::kSlowSource},
       {"CapacitorProto", BottleneckType::kSlowSource},
       {"LMDB", BottleneckType::kSlowSource},
       {"ExternalDataset", BottleneckType::kSlowSource},
       {"PearModel", BottleneckType::kSlowSource},
       {"FixedLengthRecordV2", BottleneckType::kSlowSource},
       // Read from local memory.
       {"FromTensor", BottleneckType::kSlowSource},
       {"TensorSlice", BottleneckType::kSlowSource},
       {"Generator", BottleneckType::kSlowSource},
       {"SyntheticDatasetOp", BottleneckType::kSlowSource},
       // Read from remote memory.
       {"GuzzlerDataGuzzlerRemoteDataset", BottleneckType::kSlowRemoteSource},
       {"ReverbDataset", BottleneckType::kSlowRemoteSource},
       {"DatasetService", BottleneckType::kSlowRemoteSource},
       {"DatasetSampleGame", BottleneckType::kSlowRemoteSource},
       {"Courier", BottleneckType::kSlowRemoteSource},
       {"ReverbEpisodeDataset", BottleneckType::kSlowRemoteSource},
       // Transformations with parallel version.
       {"Map", BottleneckType::kSlowTransformationWithParallelVersion},
       {"Interleave", BottleneckType::kSlowTransformationWithParallelVersion},
       // Transformations without parallel version.
       {"Filter", BottleneckType::kSlowTransformationWithoutParallelVersion},
       {"Batch", BottleneckType::kSlowTransformationWithoutParallelVersion},
       {"Unbatch", BottleneckType::kSlowTransformationWithoutParallelVersion}});
  if (auto type =
          gtl::FindOrNull(*kBottleneckTypeMap, bottleneck_iterator_name)) {
    return *type;
  }
  return BottleneckType::kOther;
}

void CombinedTfDataStatsBuilder::Add(absl::string_view host_name,
                                     XPlane* host_plane) {
  TfDataStats& tf_data_stats =
      (*combined_tf_data_stats_
            ->mutable_tf_data_stats())[std::string(host_name)];
  EventForest event_forest;
  event_forest.AddPlanes(CreateTfXPlaneVisitor, {host_plane});
  event_forest.ConnectEvents();
  event_forest.ConnectTfDataEvents();
  absl::flat_hash_set<int64> device_input_pipeline_ids;
  absl::flat_hash_map<int64, std::vector<EventNode*>> root_iterator_event_map;
  ProcessEventForest(event_forest, &device_input_pipeline_ids,
                     &root_iterator_event_map, &tf_data_stats);
  ProcessInputPipelines(device_input_pipeline_ids, &root_iterator_event_map,
                        &tf_data_stats);
}

void CombinedTfDataStatsBuilder::Finalize() {
  TfDataBottleneckAnalysis* bottleneck_analysis =
      combined_tf_data_stats_->mutable_bottleneck_analysis();
  for (const auto& host_name_and_tf_data_stats :
       combined_tf_data_stats_->tf_data_stats()) {
    SetBottleneckAnalysis(host_name_and_tf_data_stats.first,
                          host_name_and_tf_data_stats.second,
                          bottleneck_analysis);
  }
  if (generate_suggestion_) SetSuggestion(bottleneck_analysis);
  SetSummary(combined_tf_data_stats_);
}

}  // namespace profiler
}  // namespace tensorflow
