#include "tensorflow/core/tfrt/fallback/cost_recorder.h"
#include "tfrt/init_tfrt_dialects.h"  // from @tf_runtime

namespace tensorflow {
namespace {

constexpr char kCostAttrName[] = "_tfrt_cost";
constexpr char kOpKeyAttrName[] = "op_key";

absl::flat_hash_map<int64_t, uint64_t> GetOpCostMap(mlir::ModuleOp op) {
  absl::flat_hash_map<int64_t, uint64_t> op_cost_map;
  op.walk([&](mlir::Operation* op) {

#define TENSORFLOW_COMPILER_MLIR_TFRT_CONSTANTS_H_

namespace tensorflow {
namespace tfrt_compiler {

// Use __ prefix to indicate this is internal attribute.
inline constexpr char kOpKeyAttrName[] = "__op_key";

}  // namespace tfrt_compiler
}  // namespace tensorflow

    return;
  }

  // Try to use the recorded cost if any.
  if (cost_recorder_ != nullptr) {
    const auto op_key_attr =
        op->getAttrOfType<mlir::IntegerAttr>(kOpKeyAttrName);
    if (op_key_attr) {
      cost_map_[op] = cost_recorder_->GetCost(op_key_attr.getInt());
      return;
    }
  }

  // These ops are cheap regardless of their input sizes.
  //