/* 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/utils/derived_timeline.h"

#include <algorithm>
#include <utility>
#include <vector>

#include "absl/algorithm/container.h"
#include "absl/container/flat_hash_map.h"
#include "absl/strings/match.h"
#include "absl/strings/str_cat.h"
#include "absl/strings/str_split.h"
#include "absl/strings/string_view.h"
#include "absl/types/optional.h"
#include "tensorflow/core/lib/gtl/map_util.h"
#include "tensorflow/core/platform/path.h"
#include "tensorflow/core/platform/types.h"
#include "tensorflow/core/profiler/protobuf/xplane.pb.h"
#include "tensorflow/core/profiler/utils/group_events.h"
#include "tensorflow/core/profiler/utils/tf_op_utils.h"
#include "tensorflow/core/profiler/utils/tf_xplane_visitor.h"
#include "tensorflow/core/profiler/utils/time_utils.h"
#include "tensorflow/core/profiler/utils/timespan.h"
#include "tensorflow/core/profiler/utils/trace_utils.h"
#include "tensorflow/core/profiler/utils/xplane_builder.h"
#include "tensorflow/core/profiler/utils/xplane_schema.h"
#include "tensorflow/core/profiler/utils/xplane_utils.h"
#include "tensorflow/core/profiler/utils/xplane_visitor.h"

namespace tensorflow {
namespace profiler {
namespace {

const absl::string_view kAnnotationDelimiter = "::";

XEvent CreateXEvent(const XEventMetadata& metadata, int64 offset_ps,
                    int64 duration_ps, int64 group_id_stat_metadata_id,
                    absl::optional<int64> group_id) {
  XEvent event;
  event.set_metadata_id(metadata.id());
  // TODO(b/150498419): Normalize with the line start time.
  event.set_offset_ps(offset_ps);
  event.set_duration_ps(duration_ps);
  if (group_id) {
    XStat* stat = event.add_stats();
    stat->set_metadata_id(group_id_stat_metadata_id);
    stat->set_int64_value(*group_id);
  }
  return event;
}

int64 GroupIdOrInvalid(absl::optional<int64> group_id) {
  if (group_id)
    return *group_id;
  else
    return DerivedXLineBuilder::kInvalidGroupId;
}

}  // namespace

void ProcessTfOpEvent(absl::string_view tf_op_full_name,
                      absl::string_view low_level_event_name, int64 offset_ps,
                      int64 duration_ps, absl::optional<int64> group_id,
                      XPlaneBuilder* plane_builder,
                      DerivedXLineBuilder* tf_name_scope_line_builder,
                      DerivedXLineBuilder* tf_op_line_builder) {
  int64 group_id_stat_metadata_id =
      plane_builder->GetOrCreateStatMetadata(GetStatTypeStr(StatType::kGroupId))
          ->id();
  TfOp tf_op = ParseTfOpFullname(tf_op_full_name);
  Category category = tf_op.category;
  int64 group_id_or_invalid = GroupIdOrInvalid(group_id);
  if (category == Category::kTensorFlow || category == Category::kJax) {
    std::vector<XEvent> name_scope_event_per_level;
    for (const auto& tf_name_scope : ParseTfNameScopes(tf_op)) {
      name_scope_event_per_level.push_back(CreateXEvent(
          *plane_builder->GetOrCreateEventMetadata(tf_name_scope), offset_ps,
          duration_ps, group_id_stat_metadata_id, group_id));
    }
    tf_name_scope_line_builder->ExpandOrAddEvents(
        name_scope_event_per_level, group_id_or_invalid, low_level_event_name);
  }
  XEventMetadata* tf_op_event_metadata =
      plane_builder->GetOrCreateEventMetadata(tf_op_full_name);
  // Set the display name to op_type so that the events of the same op_type have
  // the same color in the trace viewer.
  tf_op_event_metadata->set_display_name(TfOpEventName(tf_op));
  tf_op_line_builder->ExpandOrAddEvent(
      CreateXEvent(*tf_op_event_metadata, offset_ps, duration_ps,
                   group_id_stat_metadata_id, group_id),
      group_id_or_invalid, low_level_event_name);
}

DerivedXLineBuilder::DerivedXLineBuilder(
    XPlaneBuilder* plane, int64 line_id, absl::string_view name,
    int64 timestamp_ns, std::vector<DerivedXLineBuilder*> dependent_lines)
    : line_(plane->GetOrCreateLine(line_id)) {
  line_.SetName(name);
  line_.SetTimestampNs(timestamp_ns);
  dependent_lines_ = std::move(dependent_lines);
}

void DerivedXLineBuilder::ExpandOrAddLevelEvent(
    const XEvent& event, int64 group_id, absl::string_view low_level_event_name,
    int level) {
  int64 offset_ps = event.offset_ps();
  int64 duration_ps = event.duration_ps();
  auto& last_event = last_event_by_level_[level];
  // If last_event is not nullptr, its offset must be less than or equal to
  // the given event's offset.
  DCHECK(!last_event || last_event->OffsetPs() <= offset_ps);
  auto& last_eventinfo = last_eventinfo_by_level_[level];
  bool merge_last_event = false;
  if (last_event && last_event->MetadataId() == event.metadata_id()) {
    // If last_event is not nullptr and metadata is same, merge the given
    // event into last_event.
    DCHECK(last_eventinfo);  // last_eventinfo must be valid as well.
    // Merges event with last_event if (1) they have the same group_id
    // and (2) low_level_event_name hasn't been seen before. If
    // low_level_event has been seen before, event and last_event are actually
    // different invocations of the same Op, and so they shouldn't be merged.
    merge_last_event =
        (group_id == last_eventinfo->group_id) &&
        !last_eventinfo->low_level_event_names.contains(low_level_event_name);
  }
  if (merge_last_event) {
    // Merge event with last_event.
    last_event->SetDurationPs((offset_ps + duration_ps) -
                              last_event->OffsetPs());
    if (!low_level_event_name.empty()) {
      // One more low_level_event_name associated with last_event.
      last_eventinfo->low_level_event_names.insert(
          std::string(low_level_event_name));
    }
  } else {
    // Otherwise, reset the last events lower than or equal to the given level.
    ResetLastEvents(level);
    // And create a new event for the given level.
    last_event = line_.AddEvent(event);
    // Also create a new XEventInfo for this level.
    last_eventinfo = XEventInfo(group_id, low_level_event_name);
  }
}

void DerivedXLineBuilder::ResetLastEvents(int level) {
  for (int i = level, end = last_event_by_level_.size(); i < end; ++i) {
    last_event_by_level_[i] = absl::nullopt;
    last_eventinfo_by_level_[i] = absl::nullopt;
  }
  if (level == 0) ResetDependentLines();
}

void DeriveEventsFromAnnotations(const SymbolResolver& symbol_resolver,
                                 const GroupMetadataMap& group_metadata_map,
                                 XPlane* device_trace, bool step_info_only) {
  // Merge and sort events by Timespan as they come from different lines.
  std::vector<XEventVisitor> events;
  uint64 start_timestamp_ns = 0;
  XPlaneVisitor device_plane = CreateTfXPlaneVisitor(device_trace);
  device_plane.ForEachLine([&](const XLineVisitor& line) {
    if (IsDerivedThreadId(line.Id())) return;  // Skip overhead line.
    start_timestamp_ns = line.TimestampNs();
    line.ForEachEvent(
        [&](const XEventVisitor& event) { events.push_back(event); });
  });
  absl::c_sort(events);

  XPlaneBuilder plane(device_trace);
  DerivedXLineBuilder tf_ops(&plane, kThreadIdTfOp, kTensorFlowOpLineName,
                             start_timestamp_ns, {});
  DerivedXLineBuilder tf_name_scope(&plane, kThreadIdTfNameScope,
                                    kTensorFlowNameScopeLineName,
                                    start_timestamp_ns, {&tf_ops});
  DerivedXLineBuilder hlo_ops(&plane, kThreadIdHloOp, kXlaOpLineName,
                              start_timestamp_ns, {});
  DerivedXLineBuilder hlo_modules(&plane, kThreadIdHloModule,
                                  kXlaModuleLineName, start_timestamp_ns,
                                  {&tf_name_scope, &hlo_ops});
  DerivedXLineBuilder steps(&plane, kThreadIdStepInfo, kStepLineName,
                            start_timestamp_ns, {&hlo_modules});
  DerivedXLineBuilder source(&plane, kThreadIdSource, kSourceLineName,
                             start_timestamp_ns, {});

  int64 group_id_stat_metadata_id =
      plane.GetOrCreateStatMetadata(GetStatTypeStr(StatType::kGroupId))->id();
  int64 step_name_stat_metadata_id =
      plane.GetOrCreateStatMetadata(GetStatTypeStr(StatType::kStepName))->id();

  // Process events in order by start time.
  for (const XEventVisitor& event : events) {
    int64 offset_ps = event.OffsetPs();
    int64 duration_ps = event.DurationPs();
    absl::string_view tf_op_full_name;
    absl::string_view hlo_module_name;
    std::vector<absl::string_view> hlo_op_names;
    absl::optional<int64> group_id;
    bool is_kernel = false;
    event.ForEachStat([&](const XStatVisitor& stat) {
      if (stat.Type() == StatType::kGroupId) {
        group_id = stat.IntValue();
      } else if (stat.Type() == StatType::kLevel0 ||  // old way to carry tf_op
                 stat.Type() == StatType::kTfOp) {
        tf_op_full_name = stat.StrOrRefValue();
      } else if (stat.Type() == StatType::kHloOp) {
        hlo_op_names =
            absl::StrSplit(stat.StrOrRefValue(), kAnnotationDelimiter);
      } else if (stat.Type() == StatType::kHloModule) {
        hlo_module_name = stat.StrOrRefValue();
      } else if (stat.Type() == StatType::kKernelDetails) {
        is_kernel = true;
      }
    });
    int64 group_id_or_invalid = GroupIdOrInvalid(group_id);
    if (group_id) {
      XEvent step_event = CreateXEvent(
          *plane.GetOrCreateEventMetadata(absl::StrCat(*group_id)), offset_ps,
          duration_ps, group_id_stat_metadata_id, group_id);
      if (auto group_metadata =
              gtl::FindOrNull(group_metadata_map, *group_id)) {
        XStat* stat = step_event.add_stats();
        stat->set_metadata_id(step_name_stat_metadata_id);
        stat->set_str_value(group_metadata->name);
      }
      steps.ExpandOrAddEvent(step_event, group_id_or_invalid);
    }

    if (step_info_only) continue;

    // For HLO/TF op lines, only use kernel events (i.e. excluding memcpy or
    // allocation events).
    if (!is_kernel) continue;

    if (!hlo_module_name.empty()) {
      hlo_modules.ExpandOrAddEvent(CreateXEvent(
          *plane.GetOrCreateEventMetadata(hlo_module_name), offset_ps,
          duration_ps, group_id_stat_metadata_id, group_id));
    }

    if (!hlo_op_names.empty()) {  // GPU kernel compiled by XLA
      DCHECK(!hlo_module_name.empty());
      std::vector<XEvent> hlo_op_event_per_level;
      for (absl::string_view hlo_op_name : hlo_op_names) {
        DCHECK(!hlo_op_name.empty());
        hlo_op_event_per_level.push_back(CreateXEvent(
            *plane.GetOrCreateEventMetadata(hlo_op_name), offset_ps,
            duration_ps, group_id_stat_metadata_id, group_id));
      }
      hlo_ops.ExpandOrAddEvents(hlo_op_event_per_level, group_id_or_invalid);
      auto symbol = symbol_resolver(hlo_module_name, hlo_op_names.back());
      if (!symbol.tf_op_name.empty()) {
        ProcessTfOpEvent(symbol.tf_op_name,
                         /*low_level_event_name=*/event.Name(), offset_ps,
                         duration_ps, group_id, &plane, &tf_name_scope,
                         &tf_ops);
      }
      if (!symbol.source_info.empty()) {
        source.ExpandOrAddEvent(CreateXEvent(
            *plane.GetOrCreateEventMetadata(symbol.source_info), offset_ps,
            duration_ps, group_id_stat_metadata_id, group_id));
      }
    } else if (!tf_op_full_name.empty()) {  // GPU kernel not compiled by XLA
      ProcessTfOpEvent(tf_op_full_name,
                       /*low_level_event_name=*/event.Name(), offset_ps,
                       duration_ps, group_id, &plane, &tf_name_scope, &tf_ops);
    }
  }
  RemoveEmptyLines(device_trace);
}

void DeriveEventsFromHostTrace(const XPlane* host_trace,
                               const GroupMetadataMap& group_metadata_map,
                               std::vector<XPlane*> device_traces) {
  struct GroupLaunchInfo {  // "Group" normally means step.
    Timespan timespan;
    int32 num_launches = 0;
    uint64 max_launch_time_ps = 0ULL;
    uint64 total_launch_time_ps = 0ULL;
  };
  typedef absl::flat_hash_map<int64 /*group_id*/, GroupLaunchInfo>
      DeviceLaunchInfo;

  int num_devices = device_traces.size();
  std::vector<DeviceLaunchInfo> per_device_launch_info(num_devices);

  XPlaneVisitor host_plane = CreateTfXPlaneVisitor(host_trace);
  host_plane.ForEachLine([&](const XLineVisitor& line) {
    if (IsDerivedThreadId(line.Id())) return;
    line.ForEachEvent([&](const XEventVisitor& event) {
      absl::optional<int64> group_id;
      absl::optional<int64> device_id;
      absl::optional<int64> correlation_id;
      // Filter out API calls for cuEventRecord/cuEventQuery/cuCtxSynchronize
      // etc for now. TODO: find a better way to filter out only the memcpy and
      // kernel launch events.
      if (absl::StartsWith(event.Name(), "cu")) return;
      event.ForEachStat([&](const XStatVisitor& stat) {
        if (stat.Type() == StatType::kGroupId) {
          group_id = stat.IntValue();
        } else if (stat.Type() == StatType::kDeviceId) {
          device_id = stat.IntOrUintValue();
        } else if (stat.Type() == StatType::kCorrelationId) {
          correlation_id = stat.IntValue();
        }
      });
      if (group_id && device_id && correlation_id && *device_id >= 0 &&
          *device_id < num_devices) {
        // This is a launch event on a known device.
        GroupLaunchInfo& group_launch_info =
            per_device_launch_info[*device_id][*group_id];
        Timespan& group_span = group_launch_info.timespan;
        Timespan event_span = event.GetTimespan();
        if (group_launch_info.num_launches) {  // Existing group.
          group_span.ExpandToInclude(event_span);
        } else {
          group_span = event_span;
        }
        ++group_launch_info.num_launches;
        group_launch_info.max_launch_time_ps = std::max(
            group_launch_info.max_launch_time_ps, event_span.duration_ps());
        group_launch_info.total_launch_time_ps += event_span.duration_ps();
      }
    });
  });

  uint64 host_plane_start = GetStartTimestampNs(*host_trace);
  for (int i = 0; i < num_devices; ++i) {
    if (per_device_launch_info[i].empty()) continue;
    uint64 device_plane_start = GetStartTimestampNs(*device_traces[i]);
    XPlaneBuilder device_plane(device_traces[i]);
    XLineBuilder launch_line =
        device_plane.GetOrCreateLine(kThreadIdKernelLaunch);
    launch_line.SetName(kKernelLaunchLineName);
    launch_line.SetTimestampNs(std::min(device_plane_start, host_plane_start));
    for (const auto& kv : per_device_launch_info[i]) {
      int64 group_id = kv.first;
      const GroupLaunchInfo& group_info = kv.second;
      if (auto group_metadata = gtl::FindOrNull(group_metadata_map, group_id)) {
        XEventBuilder device_event =
            launch_line.AddEvent(*device_plane.GetOrCreateEventMetadata(
                absl::StrCat("Launch Stats for ", group_metadata->name)));
        device_event.SetTimestampNs(
            host_plane_start + PicosToNanos(group_info.timespan.begin_ps()));
        device_event.SetDurationPs(group_info.timespan.duration_ps());
        device_event.AddStatValue(*device_plane.GetOrCreateStatMetadata(
                                      GetStatTypeStr(StatType::kGroupId)),
                                  group_id);
        device_event.AddStatValue(
            *device_plane.GetOrCreateStatMetadata("num_launches"),
            group_info.num_launches);
        device_event.AddStatValue(
            *device_plane.GetOrCreateStatMetadata("max_launch_time_us"),
            PicosToMicros(group_info.max_launch_time_ps));
        device_event.AddStatValue(
            *device_plane.GetOrCreateStatMetadata("avg_launch_time_us"),
            PicosToMicros(group_info.total_launch_time_ps /
                          group_info.num_launches));
      }
    }
  }
}

void GenerateDerivedTimeLines(const GroupMetadataMap& group_metadata_map,
                              XSpace* space, bool step_info_only) {
  // TODO(profiler): Once we capture HLO protos for xla/gpu, we should use that
  // to look up tensorflow op name from hlo_module/hlo_op.
  auto dummy_symbol_resolver = [](absl::string_view hlo_module,
                                  absl::string_view hlo_op) {
    return tensorflow::profiler::Symbol();
  };
  std::vector<XPlane*> device_traces =
      FindMutablePlanesWithPrefix(space, kGpuPlanePrefix);
  for (XPlane* plane : device_traces) {
    DeriveEventsFromAnnotations(dummy_symbol_resolver, group_metadata_map,
                                plane, step_info_only);
  }
}

}  // namespace profiler
}  // namespace tensorflow
