#include "tensorflow/compiler/mlir/quantization/tensorflow/passes/lift_quantizable_spots_as_functions_drq.inc"

void LiftQuantizableSpotsAsFunctionsDRQPass::runOnOperation() {
  MLIRContext* ctx = &getContext();
  RewritePatternSet patterns(ctx);
  ModuleOp module = getOperation();

  populateWithGenerated(patterns);
  patterns.add<CheckQuantizableOps>(ctx, quantization_method_, target_opset_,

  }
};

}  // namespace

void NchwConvolutionToNhwcPass::runOnOperation() {
  func::FuncOp func_op = getOperation();
  MLIRContext& ctx = getContext();

  RewritePatternSet patterns(&ctx);
  patterns.add<RewriteNchwConvolutionToNhwc>(&ctx);

  if (failed(applyPatternsAndFoldGreedily(func_op, std::move(patterns)))) {
    func_op.emitError() << "Failed to run NchwConvolutionToNhwcPass.";

    op->setLoc(unit_loc);

    return success();
  }
};

void AddQuantizationUnitLocPass::runOnOperation() {
  MLIRContext* ctx = &getContext();
  RewritePatternSet patterns(ctx);
  func::FuncOp func = getOperation();

  patterns.add<AddQuantizationUnitLoc>(ctx);
  if (failed(applyPatternsAndFoldGreedily(func, std::move(patterns)))) {

  void runOnOperation() override {
    auto module = getOperation();

    if (!checkWhetherGraphHasValidStaticLookupTables(module)) {
      return;
    }

    RewritePatternSet patterns(&getContext());
    patterns
        .add<LegalizeHashTableOpPattern, LegalizeHashTableFindOpPattern,
             LegalizeHashTableImportOpPattern, LegalizeHashTableSizeOpPattern>(
            &getContext());

                          std::move(*quantization_method_txtpb)));
    }
  }
  return success();
}

void LiftQuantizableSpotsAsFunctionsPass::runOnOperation() {
  MLIRContext* ctx = &getContext();
  RewritePatternSet patterns(ctx);
  ModuleOp module_op = getOperation();

  simple_patterns::populateWithGenerated(patterns);
  fusion_patterns::populateWithGenerated(patterns);

  // LSTM's restrict_scale requirement should be handled before converting stats
  // to Q-DQ ops. The pattern is applied for non-PTQ case to make op ordering
  // consistent. Otherwise some FileCheck tests would fail.
  RewritePatternSet patterns_1(&getContext());
  if (quant_specs_.post_training_quantization) {
    patterns_1.add<PrepareLstmOutputScale<LSTMOp>>(ctx);
    patterns_1.add<PrepareLstmOutputScale<UnidirectionalSequenceLSTMOp>>(ctx);
  }

// Canonicalize operations in functions.
struct TensorFlowOptimizePass
    : public impl::TensorFlowOptimizePassBase<TensorFlowOptimizePass> {
  LogicalResult initialize(MLIRContext *context) override {
    RewritePatternSet pattern_list(context);
    populateWithGenerated(pattern_list);
    pattern_list.add<SimplifyBroadcastReshape>(context);
    patterns = std::move(pattern_list);
    return success();
  }

// options and allow modular control of quantization specs.
void QuantizeWeightPass::runOnOperation() {
  func::FuncOp func = getOperation();
  MLIRContext* ctx = func.getContext();
  RewritePatternSet patterns(ctx);

  patterns.add<QuantizeWeight>(ctx, quantization_component_spec_);

  FrozenRewritePatternSet frozen_patterns(std::move(patterns));

 private:
  void runOnOperation() override;
};

void ConvertFuncToBfloat16Pass::runOnOperation() {
  func::FuncOp func_op = getOperation();
  MLIRContext* context = func_op.getContext();
  RewritePatternSet patterns(context);

  BFloat16TypeConverter converter;
  patterns.add<BFloat16TypePattern, BitcastConvertOpPattern>(converter,
                                                             context);

    // TODO(b/299514833): Remove TensorFlowDialect usage.
    registry.insert<shape::ShapeDialect, mlir::TF::TensorFlowDialect>();
  }

  void runOnOperation() override {
    RewritePatternSet patterns(&getContext());
    patterns.add<UnfuseBatchNormTrainingPattern>(&getContext());
    patterns.add<UnfuseBatchNormInferencePattern>(&getContext());
    if (failed(applyPatternsAndFoldGreedily(getOperation(),