#include <string>
#include <vector>

#include "mlir/Dialect/Quant/QuantTypes.h"  // from @llvm-project
#include "mlir/IR/OperationSupport.h"  // from @llvm-project
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops_a_m.h"

namespace mlir {
namespace TFL {

// Moves the TF operations out from the tfl.TFCustomOps wrappers inside the

#include "mlir/IR/Value.h"  // from @llvm-project
#include "mlir/Transforms/DialectConversion.h"  // from @llvm-project
#include "stablehlo/dialect/StablehloOps.h"  // from @stablehlo
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"

namespace mlir {
namespace odml {

namespace {

// Convert op to stablehlo.custom_call
//   "tf.ResizeBilinear"(%341, %138) {
//      align_corners = false, device = "", half_pixel_centers = true}

#include "tensorflow/compiler/mlir/lite/experimental/tac/hardwares/gpu_hardware.h"
#include "tensorflow/compiler/mlir/lite/ir/tfl_ops.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"

namespace mlir {
namespace TFL {
namespace tac {
namespace {

struct DeviceTransformGPUPass
    : public mlir::PassWrapper<DeviceTransformGPUPass,

include "mlir/IR/OpBase.td"
include "mlir/Dialect/Func/IR/FuncOps.td"
include "mlir/Dialect/Arith/IR/ArithOps.td"
include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.td"
include "tensorflow/compiler/mlir/quantization/common/attrs_and_constraints.td"
include "tensorflow/compiler/mlir/quantization/common/lift_as_function_call.td"

      for (OpResult per_replica_result : replicated_output->getResults()) {
        auto [partitioned_output, id_ops] =
            GetSingleUserOfType<mlir::TF::TPUPartitionedOutputV2Op>(
                per_replica_result);
        if (partitioned_output) {
          erase_list.insert(erase_list.end(), id_ops.begin(), id_ops.end());
          partitioned_outputs[cluster_result_id].emplace_back(
              partitioned_output);

    //    `reshaped_permuted` according to `crops` to produce the output of
    //    shape:
    //      [batch / prod(block_shape),
    //
    //       input_shape[1] * block_shape[0] - crops[0,0] - crops[0,1],
    //       ...,
    //       input_shape[M] * block_shape[M-1] - crops[M-1,0] - crops[M-1,1],
    //
    //       input_shape[M+1], ..., input_shape[N-1]]

#include "mlir/Support/LogicalResult.h"  // from @llvm-project
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"  // from @llvm-project
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
#include "tensorflow/core/util/matmul_bcast.h"

namespace mlir {
namespace TF {

namespace {

// Replace TF BatchMatMul by TF Einsum op
template <typename BatchMatMulOpType>

#include "mlir/Transforms/InliningUtils.h"  // from @llvm-project
#include "mlir/Transforms/Passes.h"  // from @llvm-project
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_executor.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
#include "tensorflow/compiler/mlir/tensorflow/transforms/passes.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/error_util.h"

#define DEBUG_TYPE "tf-executor-tpu-v1-island-inlining"

namespace mlir {

#include "mlir/IR/UseDefLists.h"  // from @llvm-project
#include "mlir/Pass/Pass.h"  // from @llvm-project
#include "mlir/Support/LLVM.h"  // from @llvm-project
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_saved_model.h"
#include "tensorflow/core/platform/path.h"

namespace mlir {
namespace tf_saved_model {
namespace {

#define GEN_PASS_DEF_FREEZEASSETSPASS

#include "mlir/IR/UseDefLists.h"  // from @llvm-project
#include "mlir/Pass/Pass.h"  // from @llvm-project
#include "mlir/Support/LLVM.h"  // from @llvm-project
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_saved_model.h"

namespace mlir {
namespace tf_saved_model {
namespace {

using mlir::TF::VarHandleOp;