Added Options Argument (#600)

tenstorrent · Sep 3, 2024 · ada3b6c · ada3b6c
1 parent 4ce9a3c
commit ada3b6c
Showing 1 changed file with 22 additions and 20 deletions.
diff --git 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<mlir::LogicalResult(const llvm::Twine &)> err_handler =
-        [](const llvm::Twine &loc) { return mlir::failure(); };
+        mlir::function_ref<mlir::LogicalResult(const llvm::Twine &)>
+            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_<std::shared_ptr<void>>(m, "SharedVoidPtr")
       .def(py::init<>())