// Unwraps XlaCallModule ops without quantizable trait that call function with
// '_from_xla_call_module' trait.
class UnwrapXlaCallModuleOpPass
    : public impl::UnwrapXlaCallModuleOpPassBase<UnwrapXlaCallModuleOpPass> {
 public:
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(UnwrapXlaCallModuleOpPass)

  explicit UnwrapXlaCallModuleOpPass() = default;

 private:
  void runOnOperation() override;
};

  let summary = "Merge quantized conv/dot_general fusion with subsequent dequantize.";
  let dependentDialects = [
    "chlo::ChloDialect",
    "mlir::stablehlo::StablehloDialect",
  ];
}

def UnwrapXlaCallModuleOpPass : Pass<"stablehlo-unwrap-xla-call-module-op", "ModuleOp"> {
  let summary = "Unwrap XlaCallModuleOps into inline functions if not used for quantizing fused patterns.";
  let dependentDialects = ["TF::TensorFlowDialect"];