[Mlir-commits] [mlir] [mlir] Python: Parse ModuleOp from file path (PR #126572)
Nikhil Kalra
llvmlistbot at llvm.org
Wed Feb 12 10:23:32 PST 2025
https://github.com/nikalra updated https://github.com/llvm/llvm-project/pull/126572
>From 42a3cc6ad3d058ed10eab39b1b3e9a2a196c0a9b Mon Sep 17 00:00:00 2001
From: Nikhil Kalra <nkalra at apple.com>
Date: Mon, 10 Feb 2025 10:10:25 -0800
Subject: [PATCH 1/5] Revert "Revert "[mlir] Python: Parse ModuleOp from file
path" (#126482)"
This reverts commit 67b7a2590f39ad9ff5413adb9af162220972833e.
---
mlir/include/mlir-c/IR.h | 4 ++++
mlir/include/mlir/Bindings/Python/Nanobind.h | 1 +
mlir/lib/Bindings/Python/IRCore.cpp | 16 +++++++++++++++-
mlir/lib/CAPI/IR/IR.cpp | 10 ++++++++++
mlir/python/mlir/_mlir_libs/_mlir/ir.pyi | 3 ++-
mlir/test/python/ir/module.py | 20 ++++++++++++++++++++
6 files changed, 52 insertions(+), 2 deletions(-)
diff --git a/mlir/include/mlir-c/IR.h b/mlir/include/mlir-c/IR.h
index 7d2fd89e8560f..14ccae650606a 100644
--- a/mlir/include/mlir-c/IR.h
+++ b/mlir/include/mlir-c/IR.h
@@ -309,6 +309,10 @@ MLIR_CAPI_EXPORTED MlirModule mlirModuleCreateEmpty(MlirLocation location);
MLIR_CAPI_EXPORTED MlirModule mlirModuleCreateParse(MlirContext context,
MlirStringRef module);
+/// Parses a module from file and transfers ownership to the caller.
+MLIR_CAPI_EXPORTED MlirModule
+mlirModuleCreateParseFromFile(MlirContext context, MlirStringRef fileName);
+
/// Gets the context that a module was created with.
MLIR_CAPI_EXPORTED MlirContext mlirModuleGetContext(MlirModule module);
diff --git a/mlir/include/mlir/Bindings/Python/Nanobind.h b/mlir/include/mlir/Bindings/Python/Nanobind.h
index ca942c83d3e2f..bc8bddf4caf7e 100644
--- a/mlir/include/mlir/Bindings/Python/Nanobind.h
+++ b/mlir/include/mlir/Bindings/Python/Nanobind.h
@@ -23,6 +23,7 @@
#endif
#include <nanobind/nanobind.h>
#include <nanobind/ndarray.h>
+#include <nanobind/stl/filesystem.h>
#include <nanobind/stl/function.h>
#include <nanobind/stl/optional.h>
#include <nanobind/stl/pair.h>
diff --git a/mlir/lib/Bindings/Python/IRCore.cpp b/mlir/lib/Bindings/Python/IRCore.cpp
index 47a85c2a486fd..2e4b6d1ce35c1 100644
--- a/mlir/lib/Bindings/Python/IRCore.cpp
+++ b/mlir/lib/Bindings/Python/IRCore.cpp
@@ -6,6 +6,7 @@
//
//===----------------------------------------------------------------------===//
+#include <filesystem>
#include <optional>
#include <utility>
@@ -299,7 +300,7 @@ struct PyAttrBuilderMap {
return *builder;
}
static void dunderSetItemNamed(const std::string &attributeKind,
- nb::callable func, bool replace) {
+ nb::callable func, bool replace) {
PyGlobals::get().registerAttributeBuilder(attributeKind, std::move(func),
replace);
}
@@ -3049,6 +3050,19 @@ void mlir::python::populateIRCore(nb::module_ &m) {
},
nb::arg("asm"), nb::arg("context").none() = nb::none(),
kModuleParseDocstring)
+ .def_static(
+ "parse",
+ [](const std::filesystem::path &path,
+ DefaultingPyMlirContext context) {
+ PyMlirContext::ErrorCapture errors(context->getRef());
+ MlirModule module = mlirModuleCreateParseFromFile(
+ context->get(), toMlirStringRef(path.string()));
+ if (mlirModuleIsNull(module))
+ throw MLIRError("Unable to parse module assembly", errors.take());
+ return PyModule::forModule(module).releaseObject();
+ },
+ nb::arg("asm"), nb::arg("context").none() = nb::none(),
+ kModuleParseDocstring)
.def_static(
"create",
[](DefaultingPyLocation loc) {
diff --git a/mlir/lib/CAPI/IR/IR.cpp b/mlir/lib/CAPI/IR/IR.cpp
index f27af0ca9a2c7..999e8cbda1295 100644
--- a/mlir/lib/CAPI/IR/IR.cpp
+++ b/mlir/lib/CAPI/IR/IR.cpp
@@ -22,6 +22,7 @@
#include "mlir/IR/Location.h"
#include "mlir/IR/Operation.h"
#include "mlir/IR/OperationSupport.h"
+#include "mlir/IR/OwningOpRef.h"
#include "mlir/IR/Types.h"
#include "mlir/IR/Value.h"
#include "mlir/IR/Verifier.h"
@@ -328,6 +329,15 @@ MlirModule mlirModuleCreateParse(MlirContext context, MlirStringRef module) {
return MlirModule{owning.release().getOperation()};
}
+MlirModule mlirModuleCreateParseFromFile(MlirContext context,
+ MlirStringRef fileName) {
+ OwningOpRef<ModuleOp> owning =
+ parseSourceFile<ModuleOp>(unwrap(fileName), unwrap(context));
+ if (!owning)
+ return MlirModule{nullptr};
+ return MlirModule{owning.release().getOperation()};
+}
+
MlirContext mlirModuleGetContext(MlirModule module) {
return wrap(unwrap(module).getContext());
}
diff --git a/mlir/python/mlir/_mlir_libs/_mlir/ir.pyi b/mlir/python/mlir/_mlir_libs/_mlir/ir.pyi
index fb7efb8cd28a5..096b87b362443 100644
--- a/mlir/python/mlir/_mlir_libs/_mlir/ir.pyi
+++ b/mlir/python/mlir/_mlir_libs/_mlir/ir.pyi
@@ -46,6 +46,7 @@ import abc
import collections
from collections.abc import Callable, Sequence
import io
+from pathlib import Path
from typing import Any, ClassVar, TypeVar, overload
__all__ = [
@@ -2123,7 +2124,7 @@ class Module:
Creates an empty module
"""
@staticmethod
- def parse(asm: str | bytes, context: Context | None = None) -> Module:
+ def parse(asm: str | bytes | Path, context: Context | None = None) -> Module:
"""
Parses a module's assembly format from a string.
diff --git a/mlir/test/python/ir/module.py b/mlir/test/python/ir/module.py
index ecafcb46af217..441916b38ee73 100644
--- a/mlir/test/python/ir/module.py
+++ b/mlir/test/python/ir/module.py
@@ -1,6 +1,8 @@
# RUN: %PYTHON %s | FileCheck %s
import gc
+from pathlib import Path
+from tempfile import NamedTemporaryFile
from mlir.ir import *
@@ -27,6 +29,24 @@ def testParseSuccess():
print(str(module))
+# Verify successful parse from file.
+# CHECK-LABEL: TEST: testParseFromFileSuccess
+# CHECK: module @successfulParse
+ at run
+def testParseFromFileSuccess():
+ ctx = Context()
+ with NamedTemporaryFile(mode="w") as tmp_file:
+ tmp_file.write(r"""module @successfulParse {}""")
+ tmp_file.flush()
+ module = Module.parse(Path(tmp_file.name), ctx)
+ assert module.context is ctx
+ print("CLEAR CONTEXT")
+ ctx = None # Ensure that module captures the context.
+ gc.collect()
+ module.operation.verify()
+ print(str(module))
+
+
# Verify parse error.
# CHECK-LABEL: TEST: testParseError
# CHECK: testParseError: <
>From 8ad7566b7f22cdc75263fd87d1c5a0429c9ebe90 Mon Sep 17 00:00:00 2001
From: Nikhil Kalra <nkalra at apple.com>
Date: Mon, 10 Feb 2025 10:11:17 -0800
Subject: [PATCH 2/5] guard on fs avail
---
mlir/lib/Bindings/Python/IRCore.cpp | 4 ++++
1 file changed, 4 insertions(+)
diff --git a/mlir/lib/Bindings/Python/IRCore.cpp b/mlir/lib/Bindings/Python/IRCore.cpp
index 2e4b6d1ce35c1..81936323631dd 100644
--- a/mlir/lib/Bindings/Python/IRCore.cpp
+++ b/mlir/lib/Bindings/Python/IRCore.cpp
@@ -6,7 +6,9 @@
//
//===----------------------------------------------------------------------===//
+#if __has_include(<filesystem>)
#include <filesystem>
+#endif
#include <optional>
#include <utility>
@@ -3050,6 +3052,7 @@ void mlir::python::populateIRCore(nb::module_ &m) {
},
nb::arg("asm"), nb::arg("context").none() = nb::none(),
kModuleParseDocstring)
+#if __has_include(<filesystem>)
.def_static(
"parse",
[](const std::filesystem::path &path,
@@ -3063,6 +3066,7 @@ void mlir::python::populateIRCore(nb::module_ &m) {
},
nb::arg("asm"), nb::arg("context").none() = nb::none(),
kModuleParseDocstring)
+#endif
.def_static(
"create",
[](DefaultingPyLocation loc) {
>From 76a4ffed59785f4debb8e4046f9b64bb45f2556c Mon Sep 17 00:00:00 2001
From: Nikhil Kalra <nkalra at apple.com>
Date: Mon, 10 Feb 2025 10:16:32 -0800
Subject: [PATCH 3/5] avail
---
mlir/include/mlir/Bindings/Python/Nanobind.h | 2 ++
1 file changed, 2 insertions(+)
diff --git a/mlir/include/mlir/Bindings/Python/Nanobind.h b/mlir/include/mlir/Bindings/Python/Nanobind.h
index bc8bddf4caf7e..48b23b57df108 100644
--- a/mlir/include/mlir/Bindings/Python/Nanobind.h
+++ b/mlir/include/mlir/Bindings/Python/Nanobind.h
@@ -23,7 +23,9 @@
#endif
#include <nanobind/nanobind.h>
#include <nanobind/ndarray.h>
+#if __has_include(<filesystem>)
#include <nanobind/stl/filesystem.h>
+#endif
#include <nanobind/stl/function.h>
#include <nanobind/stl/optional.h>
#include <nanobind/stl/pair.h>
>From 6270c8e8489d7929814ee08583f2501ea7147d72 Mon Sep 17 00:00:00 2001
From: Nikhil Kalra <nkalra at apple.com>
Date: Wed, 12 Feb 2025 09:37:22 -0800
Subject: [PATCH 4/5] create new api instead of overloadign
---
mlir/include/mlir/Bindings/Python/Nanobind.h | 3 ---
mlir/lib/Bindings/Python/IRCore.cpp | 14 ++++----------
mlir/python/mlir/_mlir_libs/_mlir/ir.pyi | 11 ++++++++++-
3 files changed, 14 insertions(+), 14 deletions(-)
diff --git a/mlir/include/mlir/Bindings/Python/Nanobind.h b/mlir/include/mlir/Bindings/Python/Nanobind.h
index 48b23b57df108..ca942c83d3e2f 100644
--- a/mlir/include/mlir/Bindings/Python/Nanobind.h
+++ b/mlir/include/mlir/Bindings/Python/Nanobind.h
@@ -23,9 +23,6 @@
#endif
#include <nanobind/nanobind.h>
#include <nanobind/ndarray.h>
-#if __has_include(<filesystem>)
-#include <nanobind/stl/filesystem.h>
-#endif
#include <nanobind/stl/function.h>
#include <nanobind/stl/optional.h>
#include <nanobind/stl/pair.h>
diff --git a/mlir/lib/Bindings/Python/IRCore.cpp b/mlir/lib/Bindings/Python/IRCore.cpp
index 81936323631dd..827db5f3eba84 100644
--- a/mlir/lib/Bindings/Python/IRCore.cpp
+++ b/mlir/lib/Bindings/Python/IRCore.cpp
@@ -6,9 +6,6 @@
//
//===----------------------------------------------------------------------===//
-#if __has_include(<filesystem>)
-#include <filesystem>
-#endif
#include <optional>
#include <utility>
@@ -3052,21 +3049,18 @@ void mlir::python::populateIRCore(nb::module_ &m) {
},
nb::arg("asm"), nb::arg("context").none() = nb::none(),
kModuleParseDocstring)
-#if __has_include(<filesystem>)
.def_static(
- "parse",
- [](const std::filesystem::path &path,
- DefaultingPyMlirContext context) {
+ "parseFile",
+ [](const std::string &path, DefaultingPyMlirContext context) {
PyMlirContext::ErrorCapture errors(context->getRef());
MlirModule module = mlirModuleCreateParseFromFile(
- context->get(), toMlirStringRef(path.string()));
+ context->get(), toMlirStringRef(path));
if (mlirModuleIsNull(module))
throw MLIRError("Unable to parse module assembly", errors.take());
return PyModule::forModule(module).releaseObject();
},
- nb::arg("asm"), nb::arg("context").none() = nb::none(),
+ nb::arg("path"), nb::arg("context").none() = nb::none(),
kModuleParseDocstring)
-#endif
.def_static(
"create",
[](DefaultingPyLocation loc) {
diff --git a/mlir/python/mlir/_mlir_libs/_mlir/ir.pyi b/mlir/python/mlir/_mlir_libs/_mlir/ir.pyi
index 096b87b362443..ab975a6954044 100644
--- a/mlir/python/mlir/_mlir_libs/_mlir/ir.pyi
+++ b/mlir/python/mlir/_mlir_libs/_mlir/ir.pyi
@@ -2124,12 +2124,21 @@ class Module:
Creates an empty module
"""
@staticmethod
- def parse(asm: str | bytes | Path, context: Context | None = None) -> Module:
+ def parse(asm: str | bytes, context: Context | None = None) -> Module:
"""
Parses a module's assembly format from a string.
Returns a new MlirModule or raises an MLIRError if the parsing fails.
+ See also: https://mlir.llvm.org/docs/LangRef/
+ """
+ @staticmethod
+ def parseFile(path: str, context: Context | None = None) -> Module:
+ """
+ Parses a module's assembly format from file.
+
+ Returns a new MlirModule or raises an MLIRError if the parsing fails.
+
See also: https://mlir.llvm.org/docs/LangRef/
"""
def _CAPICreate(self) -> Any: ...
>From 0211208c19750cc6eb5fe6f774d7a547381ab74e Mon Sep 17 00:00:00 2001
From: Nikhil Kalra <nkalra at apple.com>
Date: Wed, 12 Feb 2025 09:38:23 -0800
Subject: [PATCH 5/5] update test
---
mlir/test/python/ir/module.py | 3 +--
1 file changed, 1 insertion(+), 2 deletions(-)
diff --git a/mlir/test/python/ir/module.py b/mlir/test/python/ir/module.py
index 441916b38ee73..6065e59fd6ed9 100644
--- a/mlir/test/python/ir/module.py
+++ b/mlir/test/python/ir/module.py
@@ -1,7 +1,6 @@
# RUN: %PYTHON %s | FileCheck %s
import gc
-from pathlib import Path
from tempfile import NamedTemporaryFile
from mlir.ir import *
@@ -38,7 +37,7 @@ def testParseFromFileSuccess():
with NamedTemporaryFile(mode="w") as tmp_file:
tmp_file.write(r"""module @successfulParse {}""")
tmp_file.flush()
- module = Module.parse(Path(tmp_file.name), ctx)
+ module = Module.parseFile(tmp_file.name, ctx)
assert module.context is ctx
print("CLEAR CONTEXT")
ctx = None # Ensure that module captures the context.
More information about the Mlir-commits
mailing list