// Change this to user-given bit width once we have custom configuration.
  options.bit_width_ = 8;

  // Insert quantization parameters for weights for ops with `weight_only_ptq`
  // attribute.
  pm.addNestedPass<func::FuncOp>(createInsertWeightParamPass());

  // PrepareQuantizePass uses SymbolTable to fetch relevant GEMM ops for
  // determining quantization attributes. This requires module-level context.

    ModuleOp module = op->getParentOfType<ModuleOp>();
    SymbolTable symbol_table(module);
    func::FuncOp float_func =
        dyn_cast<func::FuncOp>(symbol_table.lookup(function_name));
    OperandRange func_args = op.getArgs();
    func::FuncOp new_float_func = float_func.clone();

    SmallVector<Value> new_float_func_args{func_args.begin(), func_args.end()};

#include "tensorflow/compiler/mlir/tensorflow/transforms/tf_device_passes.h.inc"

struct XlaRewritePass : public impl::XlaRewritePassBase<XlaRewritePass> {
  void runOnOperation() override;
};

void MoveResourceArgsToEnd(func::FuncOp callee) {
  llvm::DenseMap<BlockArgument, BlockArgument> mapping;
  unsigned num_params = callee.getNumArguments();
  llvm::BitVector removed_params(num_params);
  // Copy the resource-type parameters to the end.

struct DeviceIndexSelector
    : public impl::DeviceIndexSelectorPassBase<DeviceIndexSelector> {
  void runOnOperation() override;
};

}  // namespace

void DeviceIndexSelector::runOnOperation() {
  func::FuncOp func = getOperation();
  // Convert all the DeviceIndex ops to constant values.
  func.getBody().walk([](TF::DeviceIndexOp op) {
    // This just selects the default in all cases where DeviceIndex feeds into

#include "tensorflow/compiler/mlir/lite/utils/utils.h"
#include "tensorflow/compiler/mlir/quantization/common/quantization_lib/quantization_utils.h"

namespace mlir {
namespace TFL {

void ConvertTFLQuantOpsToMlirQuantOps(func::FuncOp func) {
  OpBuilder b(func);
  func.walk([&](Operation* op) {
    b.setInsertionPoint(op);
    if (auto dq = llvm::dyn_cast<DequantizeOp>(op)) {
      auto dcast = b.create<quantfork::DequantizeCastOp>(

  void EnqueueCallers(func::FuncOp fn);

  // Returns the function at the front of the queue.
  func::FuncOp front() { return queue_.front(); }

  // Returns whether work queue is empty.
  bool EmptyQueue() const { return queue_.empty(); }

  // Returns function from the front of the work queue.
  func::FuncOp pop_front() {
    func::FuncOp ret = queue_.front();
    queue_.pop();

absl::Status SetupArguments(mlir::ModuleOp module,
                            std::vector<TensorShape>& arg_shapes,
                            tpu::TPUCompileMetadataProto& metadata_proto) {
  auto main_fn = module.lookupSymbol<mlir::func::FuncOp>(kEntryFuncName);
  if (!main_fn) {
    return absl::InternalError("Could not find main function in MLIR Module.");
  }

  mlir::FunctionType func_type = main_fn.getFunctionType();

  if (old_val.getType() == new_val.getType()) {
    old_val.replaceAllUsesWith(new_val);
    return;
  }
  Operation* old_op = old_val.getDefiningOp();
  Operation* terminator_op =
      old_op->getParentOfType<func::FuncOp>().front().getTerminator();
  llvm::SmallPtrSet<mlir::Operation*, 1> exceptions = {terminator_op};
  old_val.replaceAllUsesExcept(new_val, exceptions);
}

struct TensorArrayStats {

        resource_alias_info_map.empty()) {
      return;
    }

    FlatSymbolRefAttr func_attr = cluster_func.getFuncAttr();
    func::FuncOp device_func =
        module.lookupSymbol<func::FuncOp>(func_attr.getValue());
    AddAliasingAttributeToDeviceFunc(device_func, resource_alias_info_map);
  });
}

}  // namespace

// Helper data structure for output parameters to `ExtractSubgraphToFunc`.
// `ExtractSubgraphToFunc` adds exactly two "new" `Operations`, a FuncOp and
// a CallOp. Pass these back to the caller for setting more specific attributes
// after graph mutation has taken place.
struct OpsAdded {
  mlir::func::FuncOp func_op;
  mlir::func::CallOp call_op;
};

// Given a `Subgraph` containing a sequence of adjacent `Operations` from