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

#include "absl/strings/str_cat.h"
#include "absl/strings/string_view.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/platform/types.h"
#include "tensorflow/core/profiler/convert/xplane_to_op_stats.h"
#include "tensorflow/core/profiler/protobuf/op_stats.pb.h"
#include "tensorflow/core/profiler/protobuf/tf_stats.pb.h"
#include "tensorflow/core/profiler/protobuf/xplane.pb.h"
#include "tensorflow/core/profiler/utils/time_utils.h"
#include "tensorflow/core/profiler/utils/xplane_builder.h"
#include "tensorflow/core/profiler/utils/xplane_schema.h"
#include "tensorflow/core/profiler/utils/xplane_test_utils.h"

namespace tensorflow {
namespace profiler {
namespace {

XEventBuilder AddTensorFlowOpEvent(std::string&& tf_op_fullname,
                                   int64 start_timestamp_ns, int64 duration_ns,
                                   bool on_device,
                                   absl::string_view kernel_name,
                                   XPlaneBuilder* plane, XLineBuilder* line) {
  absl::string_view name = on_device ? kernel_name : tf_op_fullname;
  XEventBuilder event = line->AddEvent(*plane->GetOrCreateEventMetadata(name));
  event.SetTimestampNs(start_timestamp_ns);
  event.SetDurationNs(duration_ns);
  if (!on_device) return event;
  event.AddStatValue(
      *plane->GetOrCreateStatMetadata("level 0"),
      *plane->GetOrCreateStatMetadata(std::move(tf_op_fullname)));
  return event;
}

void AddTensorFlowOpEventWithKernelDetails(std::string&& tf_op_fullname,
                                           int64 start_timestamp_ns,
                                           int64 duration_ns, bool on_device,
                                           absl::string_view kernel_name,
                                           absl::string_view kernel_details,
                                           XPlaneBuilder* plane,
                                           XLineBuilder* line) {
  XEventBuilder event =
      AddTensorFlowOpEvent(std::move(tf_op_fullname), start_timestamp_ns,
                           duration_ns, on_device, kernel_name, plane, line);
  if (!on_device) return;
  event.ParseAndAddStatValue(*plane->GetOrCreateStatMetadata("kernel_details"),
                             kernel_details);
}

TEST(OpStatsToTfStats, GpuTfStats) {
  // TfOp1 has kernel1 and kernel2; TfOp2 has kernel3;
  // TfOp3 has kernel4 and kernel5 and is TensorCore eligible.
  static constexpr char kTfOp1[] = "TfOp1";
  static constexpr char kTfOp2[] = "TfOp2";
  static constexpr char kTfOp3[] = "Conv2D";
  static constexpr char kKernel1[] = "kernel1";
  static constexpr char kKernel2[] = "kernel2";
  static constexpr char kKernel3[] = "kernel3";
  // Kernel4 is a kernel using TensorCore
  static constexpr char kKernel4[] = "volta_fp16_s884gemm";
  static constexpr char kKernel5[] = "kernel5";
  constexpr int64 kKernel1StartNs = 100000;
  constexpr int64 kKernel1DurationNs = 8000;
  constexpr int64 kKernel2StartNs = 110000;
  constexpr int64 kKernel2DurationNs = 10000;
  constexpr int64 kKernel3StartNs = 120000;
  constexpr int64 kKernel3DurationNs = 10000;
  constexpr int64 kKernel4StartNs = 130000;
  constexpr int64 kKernel4DurationNs = 10000;
  constexpr int64 kKernel5StartNs = 150000;
  constexpr int64 kKernel5DurationNs = 10000;

  // Mock kernel details for both kernel4 and kernel5.
  const std::string kKernelDetails = R"MULTI(registers_per_thread:32
static_shared_memory_usage:0
dynamic_shared_memory_usage:16384
grid_x:2
grid_y:1
grid_z:1
block_x:32
block_y:1
block_z:1)MULTI";

  XSpace space;
  XPlaneBuilder device_plane(
      GetOrCreateGpuXPlane(&space, /*device_ordinal=*/0));
  XLineBuilder stream1 = device_plane.GetOrCreateLine(/*line_id=*/10);
  AddTensorFlowOpEvent(absl::StrCat(kTfOp1, ":", kTfOp1), kKernel1StartNs,
                       kKernel1DurationNs, /*on_device=*/true, kKernel1,
                       &device_plane, &stream1);
  AddTensorFlowOpEvent(absl::StrCat(kTfOp1, ":", kTfOp1), kKernel2StartNs,
                       kKernel2DurationNs, /*on_device=*/true, kKernel2,
                       &device_plane, &stream1);
  XLineBuilder stream2 = device_plane.GetOrCreateLine(/*line_id=*/20);
  AddTensorFlowOpEvent(absl::StrCat(kTfOp1, ":", kTfOp1), kKernel1StartNs,
                       kKernel1DurationNs, /*on_device=*/true, kKernel1,
                       &device_plane, &stream2);
  AddTensorFlowOpEvent(absl::StrCat(kTfOp1, ":", kTfOp1), kKernel2StartNs,
                       kKernel2DurationNs, /*on_device=*/true, kKernel2,
                       &device_plane, &stream2);
  AddTensorFlowOpEvent(absl::StrCat(kTfOp2, ":", kTfOp2), kKernel3StartNs,
                       kKernel3DurationNs, /*on_device=*/true, kKernel3,
                       &device_plane, &stream2);
  AddTensorFlowOpEventWithKernelDetails(
      absl::StrCat(kTfOp3, ":", kTfOp3), kKernel4StartNs, kKernel4DurationNs,
      /*on_device=*/true, kKernel4, kKernelDetails, &device_plane, &stream2);
  AddTensorFlowOpEventWithKernelDetails(
      absl::StrCat(kTfOp3, ":", kTfOp3), kKernel5StartNs, kKernel5DurationNs,
      /*on_device=*/true, kKernel5, kKernelDetails, &device_plane, &stream2);

  OpStatsOptions options;
  options.generate_kernel_stats_db = true;
  options.generate_op_metrics_db = true;
  const OpStats op_stats = ConvertXSpaceToOpStats(space, options);
  const TfStatsDatabase tf_stats = ConvertOpStatsToTfStats(op_stats);

  EXPECT_EQ(tf_stats.device_type(), op_stats.run_environment().device_type());

  // TfOp1, TfOp3, TfOp2, Idle
  EXPECT_EQ(4, tf_stats.with_idle().tf_stats_record_size());

  const TfStatsRecord& record_0 = tf_stats.with_idle().tf_stats_record(0);
  EXPECT_EQ(kTfOp1, record_0.op_name());
  EXPECT_EQ(kTfOp1, record_0.op_type());
  EXPECT_EQ(2, record_0.occurrences());
  EXPECT_EQ(NanosToMicros(kKernel1DurationNs) * 2 +
                NanosToMicros(kKernel2DurationNs) * 2,
            record_0.total_self_time_in_us());

  const TfStatsRecord& record_1 = tf_stats.with_idle().tf_stats_record(1);
  EXPECT_EQ(kTfOp3, record_1.op_name());
  EXPECT_EQ(kTfOp3, record_1.op_type());
  EXPECT_EQ(1, record_1.occurrences());
  EXPECT_EQ(
      NanosToMicros(kKernel4DurationNs) + NanosToMicros(kKernel5DurationNs),
      record_1.total_self_time_in_us());
  // GPU TensorCore utilization is 0.5 because kernel4 is using TensorCore and
  // kernel5 is not using TensorCore, and they have the same duration.
  EXPECT_DOUBLE_EQ(0.5, record_1.gpu_tensorcore_utilization());

  const TfStatsRecord& record_2 = tf_stats.with_idle().tf_stats_record(2);
  EXPECT_EQ(kTfOp2, record_2.op_name());
  EXPECT_EQ(kTfOp2, record_2.op_type());
  EXPECT_EQ(1, record_2.occurrences());
  EXPECT_EQ(NanosToMicros(kKernel3DurationNs),
            record_2.total_self_time_in_us());
}

}  // namespace
}  // namespace profiler
}  // namespace tensorflow
