/* Copyright 2025 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_PERFORMANCE_INFO_WRAPPER_H_
#define XPROF_UTILS_PERFORMANCE_INFO_WRAPPER_H_

#include <cstdint>
#include <memory>
#include <optional>
#include <vector>

#include "xla/hlo/ir/hlo_instruction.h"
#include "xla/hlo/ir/hlo_opcode.h"
#include "tsl/platform/protobuf.h"
#include "plugin/xprof/protobuf/op_metrics.pb.h"
#include "xprof/utils/cost_utils.h"
#include "xprof/utils/hlo_cost_analysis_wrapper.h"

namespace tensorflow {
namespace profiler {

// Wrapper class for handling PerformanceInfo generated by HloCostAnalysis.
class PerformanceInfoWrapper {
 public:
  using PerfInfoType = tensorflow::profiler::PerformanceInfo;
  static std::unique_ptr<PerformanceInfoWrapper> Create(
      const tensorflow::profiler::HloCostAnalysisWrapper* hlo_cost_analysis,
      const xla::HloInstruction* hlo_instruction);

  static std::unique_ptr<PerformanceInfoWrapper> Create(
      std::unique_ptr<PerfInfoType> performance_info);

  int64_t BytesAccessed() const {
    return ValidHloCost(performance_info_->bytes_accessed());
  }

  // Number of flops required for the hlo_instruction.
  int64_t ModelFlops() const {
    return ValidHloCost(performance_info_->flops());
  }
  // Number of flops normalized to the peak capacity of the device.
  int64_t DeviceFlops() const;

  // The bit width of the computation. Example, If inputs are bfloat16, the
  // computation bit width would be 16.
  int64_t ComputationalPrimitiveBitwidth() const;

  // Get the Input bit widths for the computation.
  std::vector<uint32_t> InputBitwidths() const;

  const tsl::protobuf::RepeatedPtrField<PerfInfoType::MemoryAccessed>&
  MemoryAccessedBreakdown() const {
    return performance_info_->memory_accessed_breakdown();
  }

  // Functions to allow PerformanceInfoWrapper to be used in place of
  // PerformanceInfo. EG in Roofline Analysis.
  int64_t bytes_accessed() const { return BytesAccessed(); }
  int64_t flops() const { return ModelFlops(); }
  const tsl::protobuf::RepeatedPtrField<PerfInfoType::MemoryAccessed>&
  memory_accessed_breakdown() const {
    return MemoryAccessedBreakdown();
  }

  // Convert XLA's memory access breakdown into TF's counterpart.
  // TODO(b/408280338) code cleanup.
  tensorflow::profiler::MemoryAccessBreakdown GetMemmoryAccessBreakdown() const;

 private:
  PerformanceInfoWrapper(std::unique_ptr<PerfInfoType> performance_info,
                         std::optional<xla::HloOpcode> opcode,
                         int64_t device_flops,
                         std::vector<uint32_t> input_bitwidths);

  const std::unique_ptr<PerfInfoType> performance_info_;
  const std::optional<xla::HloOpcode> opcode_;
  const int64_t device_flops_;
  const std::vector<uint32_t> input_bitwidths_;
};

}  // namespace profiler
}  // namespace tensorflow

#endif  // XPROF_UTILS_PERFORMANCE_INFO_WRAPPER_H_
