/* Copyright 2019 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.
==============================================================================*/

#ifndef XPROF_UTILS_OP_METRICS_DB_UTILS_H_
#define XPROF_UTILS_OP_METRICS_DB_UTILS_H_

#include <algorithm>
#include <cstdint>
#include <optional>
#include <string>
#include <vector>

#include "absl/container/flat_hash_map.h"
#include "absl/strings/string_view.h"
#include "xla/tsl/platform/macros.h"
#include "xla/tsl/platform/types.h"
#include "xla/tsl/profiler/utils/xplane_visitor.h"
#include "plugin/xprof/protobuf/op_metrics.pb.h"

namespace tensorflow {
namespace profiler {

// The name of OpMetrics to represent the idle time.
TF_CONST_INIT extern const absl::string_view kIdle;
// The core index to add to sparse core index in op metrics.
TF_CONST_INIT extern const uint32_t kSparseCoreIndexStart;

// Helps build an op metrics database (borrowed).
// Enables fast lookup of existing ops and prevents the creation of duplicate
// ops. It is the user's responsibility to ensure an op metrics database
// outlives its builder, and that no ops are added to the database outside of
// the builder.
class OpMetricsDbBuilder {
 public:
  // Create with a borrowed op database.
  // REQUIRED: The op database must be empty.
  explicit OpMetricsDbBuilder(OpMetricsDb* db);

 protected:
  // Looks up the given OP name. If it is already in the database,
  // return its OpMetrics; otherwise, insert a new one.
  OpMetrics* LookupOrInsertNewOpMetrics(uint64_t hlo_module_id,
                                        absl::string_view name);

  OpMetricsDb* db() { return db_; }

 private:
  // Map op (hlo_module_id, name) to the corresponding metrics in the op
  // database.
  absl::flat_hash_map<uint64_t /*hlo_module_id*/,
                      absl::flat_hash_map<std::string /*name*/, OpMetrics*>>
      op_metrics_map_;

  // The op database.
  OpMetricsDb* db_;
};

// Helps build an op metrics database (borrowed) from XEvents,
class XEventsOpMetricsDbBuilder {
 public:
  struct OpKey {
    std::optional<uint64_t> program_id;
    std::optional<uint64_t> symbol_id;
  };
  // DEPRECATED: Use the OpKey version below.
  // Add OpMetric from XEventVisitor.
  void AddOpMetric(const tsl::profiler::XEventVisitor& xevent);

  // Add an OpMetric to the builder based on the provided key.
  void AddOpMetric(const OpMetrics& op_metrics, const OpKey& key);

  // Finalize OpMetricDb and add total time and Idle op.
  OpMetricsDb Finalize(uint64_t total_time);

  // Finalize OpMetricDb, but the total time is unknown at the moment, So ignore
  // the total time and Idle Op and will be handled by the caller.
  OpMetricsDb Finalize();

 private:
  using OpMetricBySymbol =
      absl::flat_hash_map</*symbol_id=*/uint64_t, OpMetrics>;
  absl::flat_hash_map</*program_id=*/uint64_t, OpMetricBySymbol>
      flat_op_metric_;
};

// Constructs an OpMetrics from the provided XEventVisitor.
OpMetrics FromXEvent(const tsl::profiler::XEventVisitor& xevent);

// Returns the OpKey for the provided XEventVisitor.
XEventsOpMetricsDbBuilder::OpKey GetOpKeyFromXEvent(
    const tsl::profiler::XEventVisitor& event);

// Sets the total time for OpMetricsDb, ensuring idle time is not negative.
inline void SetTotalTimePs(OpMetricsDb& db, uint64_t total_time_ps) {
  db.set_total_time_ps(std::max(db.total_op_time_ps(), total_time_ps));
}

// Returns the total time in OpMetricsDb, optionally excluding the idle time.
inline uint64_t TotalTimePs(const OpMetricsDb& db, bool exclude_idle = false) {
  return exclude_idle ? db.total_op_time_ps() : db.total_time_ps();
}

// Returns the ratio of time that is idle (no op execution) over total time.
double IdleTimeRatio(const OpMetricsDb& db);

// Returns the idle time in picoseconds.
uint64_t IdleTimePs(const OpMetricsDb& db);

// Populates an OpMetrics record representing idle time, i.e., the amount of
// time spent without any op execution.
void SetIdleOp(uint64_t idle_time_ps, OpMetrics& metrics);

// Adds an OpMetrics record representing idle time, i.e., the amount of time
// spent without any op execution.
// REQUIRED: All ops must have been added to the database and the total time
// must have been set.
void AddIdleOp(OpMetricsDb& db);

// Returns true if the given metrics represents idle time.
inline bool IsIdleOp(const OpMetrics& metrics) {
  return metrics.category() == kIdle;
}

// Returns the time spent in children (nested) ops.
inline uint64_t ChildrenTimePs(const OpMetrics& metrics) {
  return metrics.time_ps() - metrics.self_time_ps();
}

// Returns the ratio of time spent sending data from the host to the device
// relative to the total time the host was active.
std::optional<double> HostInfeedEnqueueRatio(const OpMetricsDb& db);

// Parses the provenance string into a vector of strings.
std::vector<std::string> ParseProvenance(absl::string_view provenance);

// Converts from the device op metrics to Tf-op metrics.
OpMetricsDb CreateTfMetricsDbFromDeviceOpMetricsDb(
    const OpMetricsDb& device_op_metrics_db, bool with_idle = true);

// Returns the root nodes of `op_metrics`, which are assumed to be nodes in a
// forest. Every node is assumed to be uniquely identified by its
// `hlo_module_id` and `name`.
std::vector<const OpMetrics*> GetRootOpMetricsFromDb(
    const OpMetricsDb& op_metrics_db);

}  // namespace profiler
}  // namespace tensorflow

#endif  // XPROF_UTILS_OP_METRICS_DB_UTILS_H_
