// for canonicalization and folding which is not offered directly by the
// shlo implementation.
class SHLOSimplifyPass : public impl::SHLOSimplifyPassBase<SHLOSimplifyPass> {
 public:
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(SHLOSimplifyPass)

  void runOnOperation() override {
    ModuleOp module = getOperation();
    RewritePatternSet patterns(&getContext());
    populateWithGenerated(patterns);

using ::stablehlo::quantization::QuantizationConfig;

class TestPreCalibrationComponentPass
    : public impl::TestPreCalibrationComponentPassBase<
          TestPreCalibrationComponentPass> {
 public:
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestPreCalibrationComponentPass)

 private:
  void runOnOperation() override;
};

void TestPreCalibrationComponentPass::runOnOperation() {
  ModuleOp module_op = getOperation();

namespace {

// Converts XlaCallModuleOps to func.call.
class XlaCallModuleToCallPass
    : public impl::XlaCallModuleToCallPassBase<XlaCallModuleToCallPass> {
 public:
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(XlaCallModuleToCallPass)

  explicit XlaCallModuleToCallPass() = default;

 private:
  void runOnOperation() override;
};

// Converts XlaCallModuleOps to func.call.

// which factors consecutive same-dialect ops into functions.
class OrderByDialectPass
    : public impl::OrderByDialectPassBase<OrderByDialectPass> {
 public:
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(OrderByDialectPass)
  void runOnOperation() override;
};

int DialectOrdering(Operation* predecessor, Operation* op) {
  return predecessor && predecessor->getName().getDialectNamespace() ==

// Define a pass class for checking if non-accepted dialect ops exist or not.
namespace {
class CheckAcceptedOpsPass
    : public PassWrapper<CheckAcceptedOpsPass, OperationPass<>> {
 public:
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(CheckAcceptedOpsPass)

  explicit CheckAcceptedOpsPass(
      const std::vector<std::string> &optional_accepted_dialects)
      : accepted_dialects_(GetAcceptedStableHLODialects()),

namespace mlir::quant {
namespace {

// Applies optimization after quantization.
class OptimizePass
    : public PassWrapper<OptimizePass, OperationPass<func::FuncOp>> {
 public:
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(OptimizePass)

  StringRef getArgument() const final {
    // This is the argument used to refer to the pass in
    // the textual format (on the commandline for example).
    return "quant-optimize";

namespace mlir {
namespace quant {
namespace {

class ConvertFakeQuantToQdqPass
    : public PassWrapper<ConvertFakeQuantToQdqPass,
                         OperationPass<func::FuncOp>> {
 public:
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(ConvertFakeQuantToQdqPass)

  StringRef getArgument() const final {
    // This is the argument used to refer to the pass in
    // the textual format (on the commandline for example).

    : public PassWrapper<MockSuccessPass, OperationPass<ModuleOp>> {
  void getDependentDialects(DialectRegistry& registry) const override {
    registry.insert<TF::TensorFlowDialect>();
  }

 public:
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(MockSuccessPass)

  explicit MockSuccessPass() = default;

 private:
  void runOnOperation() override {
    getOperation().walk([](Operation* nestedOp) {
      nestedOp->emitError()

// the XlaCallModuleSerialization.
class RestoreFunctionNamePass
    : public impl::RestoreFunctionNamePassBase<RestoreFunctionNamePass> {
 public:
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(RestoreFunctionNamePass)

  explicit RestoreFunctionNamePass() = default;

  void runOnOperation() override;
};

void RestoreFunctionNameFromXlaCallModuleOp(TF::XlaCallModuleOp& call_op,

// arguments to functions, to TFLite variables.
class LegalizeVariablesPass
    : public impl::LegalizeVariablesPassBase<LegalizeVariablesPass> {
 public:
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(LegalizeVariablesPass)

  void runOnOperation() override {
    auto module = getOperation();
    // If TFLite variable legalization is not allowed, then we skip this pass.