// `tf.backend_config` is a DictionaryAttr, JAX2TF sets the value of its
// i64 attribute `called_index` to the TF function's name.
constexpr llvm::StringRef kTfBackendConfigAttrName = "tf.backend_config";
constexpr llvm::StringRef kCalledIndexAttrName = "called_index";
constexpr llvm::StringRef kCalledFuncAttrName = "called_func";

// Converts `called_func` attributes in custom call ops back to `called_index`.

#include "tensorflow/compiler/mlir/tensorflow/utils/session_utils.h"

#include <string>
#include <vector>

#include "absl/status/status.h"
#include "llvm/ADT/SmallSet.h"
#include "llvm/ADT/StringRef.h"
#include "mlir/Support/LLVM.h"  // from @llvm-project
#include "tensorflow/compiler/mlir/utils/string_container_utils.h"
#include "tensorflow/core/common_runtime/device_mgr.h"
#include "tensorflow/core/framework/device.h"

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

#include "llvm/ADT/StringRef.h"
#include "mlir/IR/Builders.h"  // from @llvm-project
#include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
#include "mlir/IR/BuiltinTypeInterfaces.h"  // from @llvm-project

#include "tensorflow/compiler/mlir/tf2xla/internal/legalize_tf_to_hlo.h"

#include <memory>
#include <string>
#include <vector>

#include "absl/log/log.h"
#include "llvm/ADT/StringRef.h"
#include "mlir/Pass/Pass.h"  // from @llvm-project
#include "tensorflow/compiler/mlir/tf2xla/api/v1/compile_tf_graph.h"
#include "tensorflow/compiler/mlir/tf2xla/internal/compilation_timer.h"

using ::tsl::protobuf::TextFormat;

// Given a `quantized_func_name` that starts with `kQuantizedFuncPrefix`,
// converts `kQuantizedFuncPrefix` to `kCompositeFuncPrefix`.
std::string GetCompositeFunctionName(const StringRef quantized_func_name) {
  return Twine(kCompositeFuncPrefix)
      .concat(quantized_func_name.rsplit(kQuantizedFuncPrefix).second)
      .str();
}

  for (size_t i = 0; i < permutation.size(); ++i) {
    reverse[permutation[i]] = i;
  }
  return reverse;
}

SmallVector<int64_t, 4> GetDataFormatPermutation(StringRef from, StringRef to) {
  if (from == "NHWC" && to == "NCHW") {
    return {0, 3, 1, 2};
  } else if (from == "NCHW" && to == "NHWC") {
    return {0, 2, 3, 1};
  } else {
    return {};
  }
}

static constexpr char kArbitraryDeviceName[] = "XLA_CPU_JIT";

namespace {
inline absl::string_view StringRefToView(llvm::StringRef ref) {
  return {ref.data(), ref.size()};
}
}  // namespace

static OwningOpRef<mlir::ModuleOp> GraphdefToMlirTranslateFunction(
    llvm::StringRef input, MLIRContext* context) {
  tensorflow::GraphdefToMlirOptions options{
      debug_info_file,        xla_compile_device_type,

    if (failed(cast_i32)) return failure();

    int32_array.push_back(*cast_i32);
  }
  return int32_array;
}

StringRef GetEntryFunctionName(TF::XlaCallModuleOp op) {
  if (!op->hasAttrOfType<FlatSymbolRefAttr>(
          TF::kStablehloEntryFunctionAttrName)) {
    return StringRef();
  }
  return op
      ->getAttrOfType<FlatSymbolRefAttr>(TF::kStablehloEntryFunctionAttrName)
      .getValue();
}

                               mlir::OperationPass<ModuleOp>> {
 public:
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(FoldConstantsToSubgraphPass)

  llvm::StringRef getArgument() const final {
    return "tfl-fold-constants-to-subgraph";
  }
  llvm::StringRef getDescription() const final {
    return "Fold constants into each subgraph.";
  }
  FoldConstantsToSubgraphPass() = default;

  if (!llvm::isa<CallSiteLoc>(attr)) return false;
  auto callsite_loc = llvm::dyn_cast<CallSiteLoc>(attr);

  if (!mlir::isa<NameLoc>(callsite_loc.getCaller())) return false;
  StringRef caller_name =
      mlir::cast<NameLoc>(callsite_loc.getCaller()).getName().strref();
  return caller_name.starts_with(kQuantizationUnitPrefix) &&
         caller_name.ends_with(kQuantizationUnitSuffix);
}