#include "tensorflow/compiler/mlir/tensorflow/ir/tf_types.h"
#include "tensorflow/compiler/mlir/tensorflow/transforms/passes.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/parallel_execute_util.h"
#include "xla/xla_data.pb.h"

namespace mlir {
namespace TFTPU {
namespace {

#define GEN_PASS_DEF_TPURESOURCEREADSWRITESPARTITIONINGPASS
#include "tensorflow/compiler/mlir/tensorflow/transforms/tf_passes.h.inc"

}

def XlaBroadcastPass : Pass<"tf-xla-broadcast", "mlir::func::FuncOp"> {
  let summary = "Moves a broadcast from host into XLA, encoded as XlaAllReduce";

  let description = [{
    This pass moves brodcasts from host TF ops into XLA. This enables use of
    the inter-device network, which is faster than the inter-host network.
    Broadcasts in XLA are encoded as XlaAllReduce. An all_reduce with all 0

        "@llvm-project//mlir:FuncDialect",
        "@llvm-project//mlir:Pass",
        "@llvm-project//mlir:ReconcileUnrealizedCasts",
        "@llvm-project//mlir:Transforms",
        "@local_xla//xla/mlir_hlo:mhlo_passes",
    ],
)

cc_library(
    name = "pre_calibration",
    srcs = ["pre_calibration.cc"],
    hdrs = ["pre_calibration.h"],
    compatible_with = get_compatible_with_portable(),

#include "tensorflow/cc/ops/functional_ops.h"
#include "tensorflow/cc/ops/standard_ops.h"
#include "tensorflow/compiler/jit/encapsulate_util.h"
#include "tensorflow/compiler/tf2xla/rearrange_function_argument.h"
#include "xla/test.h"
#include "tensorflow/core/common_runtime/device_factory.h"
#include "tensorflow/core/common_runtime/function.h"
#include "tensorflow/core/framework/common_shape_fns.h"
#include "tensorflow/core/framework/function.h"

// similar manner.
//
// This pass may be used to avoid placing shape-determining constants in the CPU
// graph and pass them as arguments to the TPU graph (via `TPUPartitionedCall`).
// If this happens, the XLA compiler cannot recognize such arguments as
// constants and may result in an error.
//
// A set of predefined ops and operand indices is used to determine whether an
// operand is a target for constant duplication.

                     "Uses TF ops that mimic quantization behavior"),
          clEnumValN(OpSet::XLA, "XLA", "Uses TF XLA ops"),
          clEnumValN(OpSet::UNIFORM_QUANTIZED, "UNIFORM_QUANTIZED",
                     "Uses TF Uniform Quantized ops"))};
};

bool QuantizePass::shouldKeepUnusedQdqPattern() {
  return target_opset_ == OpSet::XLA &&
         (quant_specs_.weight_only_quantization ||
          quant_specs_.weight_quantization);

  func.return
}

// CHECK-LABEL: func @empty_func() {
func.func @empty_func() {
  func.return
}

// -----

// Tests with a block argument inputs/outputs with no xla sharding op attached
// gets default maximal(0) sharding configuration.
// CHECK-LABEL: func @check_default_sharding_for_block_arg_inputs_outputs
func.func @check_default_sharding_for_block_arg_inputs_outputs(%arg0: tensor<*xi32>) {

#include "tensorflow/compiler/mlir/tosa/tf_tfl_passes.h"
#include "tensorflow/compiler/mlir/tosa/tfl_passes.h"
#include "tensorflow/compiler/mlir/tosa/transforms/passes.h"
#include "xla/mlir/framework/transforms/passes.h"
#include "xla/mlir_hlo/mhlo/transforms/passes.h"
#include "xla/status_macros.h"
#include "tensorflow/core/common_runtime/eager/context.h"
#include "tensorflow/core/common_runtime/function_body.h"

#include "tensorflow/compiler/mlir/tf2xla/api/v1/compile_mlir_util.h"
#include "tensorflow/compiler/mlir/tf2xla/transforms/passes.h"
#include "xla/mlir/framework/transforms/passes.h"
#include "xla/mlir_hlo/mhlo/IR/register.h"
#include "xla/mlir_hlo/mhlo/transforms/passes.h"
#include "tensorflow/core/platform/errors.h"
#include "tensorflow/core/platform/status.h"
#include "tensorflow/core/platform/statusor.h"

// variables instead).
std::unique_ptr<OperationPass<ModuleOp>>
CreateRemoveVariableInitializationByConstPass();

// Creates a pass that converts Tensorflow Xla ops to non-Xla ops.
std::unique_ptr<OperationPass<func::FuncOp>> CreateConvertTfXlaOpToTfOpPass();

// Creates a pass that converts TPU models for CPU by removing TPU related ops