[Mlir-commits] [mlir] [mlir] Python: Parse ModuleOp from file path (PR #126572)
Nikhil Kalra
llvmlistbot at llvm.org
Mon Feb 10 10:44:27 PST 2025
https://github.com/nikalra created https://github.com/llvm/llvm-project/pull/126572
For extremely large models, it may be inefficient to load the model into memory in Python prior to passing it to the MLIR C APIs for deserialization. This change adds an API to parse a ModuleOp directly from a file path.
Re-lands [4e14b8a](https://github.com/llvm/llvm-project/commit/4e14b8afb44af58ab7073bb8c0b52875599b0ae1).
>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/3] 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 7d2fd89e8560fc9..14ccae650606af8 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 ca942c83d3e2fad..bc8bddf4caf7e77 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 47a85c2a486fd46..2e4b6d1ce35c1b6 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 f27af0ca9a2c78b..999e8cbda1295a1 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 fb7efb8cd28a5eb..096b87b36244368 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 ecafcb46af2175d..441916b38ee73bb 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/3] 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 2e4b6d1ce35c1b6..81936323631dd7d 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/3] 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 bc8bddf4caf7e77..48b23b57df10841 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>
More information about the Mlir-commits
mailing list