diff --git a/python/Passes.cpp b/python/Passes.cpp index d599b44e3..4fa3b23bf 100644 --- a/python/Passes.cpp +++ b/python/Passes.cpp @@ -17,31 +17,33 @@ void populatePassesModule(py::module &m) { mlir::tt::registerAllPasses(); mlir::registerAllTranslations(); - m.def("ttir_to_ttnn_backend_pipeline", [](MlirModule module) { - mlir::Operation *moduleOp = unwrap(mlirModuleGetOperation(module)); - mlir::PassManager pm(moduleOp->getName()); - - mlir::DialectRegistry registry; - mlir::tt::registerAllDialects(registry); - mlir::MLIRContext *ctx = unwrap(mlirModuleGetContext(module)); - ctx->appendDialectRegistry(registry); + m.def( + "ttir_to_ttnn_backend_pipeline", + [](MlirModule module, std::string options = "") { + mlir::Operation *moduleOp = unwrap(mlirModuleGetOperation(module)); + mlir::PassManager pm(moduleOp->getName()); - const auto pipeline = - mlir::PassPipelineInfo::lookup("ttir-to-ttnn-backend-pipeline"); + mlir::DialectRegistry registry; + mlir::tt::registerAllDialects(registry); + mlir::MLIRContext *ctx = unwrap(mlirModuleGetContext(module)); + ctx->appendDialectRegistry(registry); - std::string options = ""; + const auto pipeline = + mlir::PassPipelineInfo::lookup("ttir-to-ttnn-backend-pipeline"); - mlir::function_ref err_handler = - [](const llvm::Twine &loc) { return mlir::failure(); }; + mlir::function_ref + err_handler = + [](const llvm::Twine &loc) { return mlir::failure(); }; - if (mlir::failed(pipeline->addToPipeline(pm, options, err_handler))) { - throw std::runtime_error("Failed to add pipeline to pass manager"); - } + if (mlir::failed(pipeline->addToPipeline(pm, options, err_handler))) { + throw std::runtime_error("Failed to add pipeline to pass manager"); + } - if (mlir::failed(pm.run(moduleOp))) { - throw std::runtime_error("Failed to run pass manager"); - } - }); + if (mlir::failed(pm.run(moduleOp))) { + throw std::runtime_error("Failed to run pass manager"); + } + }, + py::arg("module"), py::arg("options") = ""); py::class_>(m, "SharedVoidPtr") .def(py::init<>())