[Mlir-commits] [mlir] [mlir][python] Import NumPy lazily in mlir.ir (PR #238232)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Mon Nov 2 23:05:41 PST 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir
Author: brandonsun (brandon-yujie-sun)
<details>
<summary>Changes</summary>
Importing mlir.ir currently imports NumPy to register four dense-elements
attribute builders, even when those builders are never used. Register the
builders unconditionally and import NumPy when a builder is called instead.
When NumPy is unavailable, importing mlir.ir still succeeds. Building one of
these attributes from Python values now raises a ModuleNotFoundError naming
NumPy, rather than failing because the attribute builder was not registered.
Missing transitive dependencies and other NumPy import failures propagate
unchanged.
Add fresh-process regression coverage for import laziness, all four builders,
missing NumPy, a missing transitive dependency, and a broken NumPy import.
Give the memref test an explicit NumPy import instead of relying on the
incidental np name exported by `from mlir.ir import *`.
Validation: check-mlir-python (108 passed, 7 unsupported). The new regression
test fails against unpatched ir.py and passes with this change. Black checks
pass for ir.py, the new test, and the changed lines in the memref test.
In 20 alternating fresh-process pairs on Linux x86-64 with Python 3.12.3,
NumPy 2.1.0, and OPENBLAS_NUM_THREADS=1, median import time was 59.8 ms before
and 27.5 ms after. Both variants used the same native MLIR build. Measurements are
specific to this configuration.
Tool assistance: Codex
---
Full diff: https://github.com/llvm/llvm-project/pull/238232.diff
3 Files Affected:
- (modified) mlir/python/mlir/ir.py (+56-37)
- (modified) mlir/test/python/dialects/memref.py (+2)
- (added) mlir/test/python/ir/lazy_numpy.py (+75)
``````````diff
diff --git a/mlir/python/mlir/ir.py b/mlir/python/mlir/ir.py
index 011ecfffd355f1..198627977a4cda 100644
--- a/mlir/python/mlir/ir.py
+++ b/mlir/python/mlir/ir.py
@@ -383,40 +383,59 @@ def _memref_type_attr(x, context):
return _typeAttr(x, context)
-try:
- import numpy as np
-
- @register_attribute_builder("F64ElementsAttr")
- def _f64ElementsAttr(x, context):
- return DenseElementsAttr.get(
- np.array(x, dtype=np.float64),
- type=F64Type.get(context=context),
- context=context,
- )
-
- @register_attribute_builder("I32ElementsAttr")
- def _i32ElementsAttr(x, context):
- return DenseElementsAttr.get(
- np.array(x, dtype=np.int32),
- type=IntegerType.get_signless(32, context=context),
- context=context,
- )
-
- @register_attribute_builder("I64ElementsAttr")
- def _i64ElementsAttr(x, context):
- return DenseElementsAttr.get(
- np.array(x, dtype=np.int64),
- type=IntegerType.get_signless(64, context=context),
- context=context,
- )
-
- @register_attribute_builder("IndexElementsAttr")
- def _indexElementsAttr(x, context):
- return DenseElementsAttr.get(
- np.array(x, dtype=np.int64),
- type=IndexType.get(context=context),
- context=context,
- )
-
-except ImportError:
- pass
+def _import_numpy():
+ """Imports NumPy on first use of a dense elements attribute builder."""
+ try:
+ import numpy
+ except ModuleNotFoundError as e:
+ if e.name != "numpy":
+ raise
+ import sys
+
+ raise ModuleNotFoundError(
+ "Building a dense elements attribute from Python values requires "
+ "NumPy but it was not installed. "
+ f"Recommend: {sys.executable} -m pip install numpy",
+ name="numpy",
+ ) from e
+ return numpy
+
+
+ at register_attribute_builder("F64ElementsAttr")
+def _f64ElementsAttr(x, context):
+ np = _import_numpy()
+ return DenseElementsAttr.get(
+ np.array(x, dtype=np.float64),
+ type=F64Type.get(context=context),
+ context=context,
+ )
+
+
+ at register_attribute_builder("I32ElementsAttr")
+def _i32ElementsAttr(x, context):
+ np = _import_numpy()
+ return DenseElementsAttr.get(
+ np.array(x, dtype=np.int32),
+ type=IntegerType.get_signless(32, context=context),
+ context=context,
+ )
+
+
+ at register_attribute_builder("I64ElementsAttr")
+def _i64ElementsAttr(x, context):
+ np = _import_numpy()
+ return DenseElementsAttr.get(
+ np.array(x, dtype=np.int64),
+ type=IntegerType.get_signless(64, context=context),
+ context=context,
+ )
+
+
+ at register_attribute_builder("IndexElementsAttr")
+def _indexElementsAttr(x, context):
+ np = _import_numpy()
+ return DenseElementsAttr.get(
+ np.array(x, dtype=np.int64),
+ type=IndexType.get(context=context),
+ context=context,
+ )
diff --git a/mlir/test/python/dialects/memref.py b/mlir/test/python/dialects/memref.py
index b91fdc367cf301..f24f587da5d4a3 100644
--- a/mlir/test/python/dialects/memref.py
+++ b/mlir/test/python/dialects/memref.py
@@ -1,5 +1,7 @@
# RUN: %PYTHON %s | FileCheck %s
+import numpy as np
+
import mlir.dialects.arith as arith
import mlir.dialects.memref as memref
import mlir.extras.types as T
diff --git a/mlir/test/python/ir/lazy_numpy.py b/mlir/test/python/ir/lazy_numpy.py
new file mode 100644
index 00000000000000..264c8cc1afa381
--- /dev/null
+++ b/mlir/test/python/ir/lazy_numpy.py
@@ -0,0 +1,75 @@
+# RUN: %PYTHON %s | FileCheck %s
+# RUN: %PYTHON %s --without-numpy | FileCheck %s --check-prefix=NO-NUMPY
+# RUN: %PYTHON %s --missing-dependency | FileCheck %s --check-prefix=DEPENDENCY
+# RUN: %PYTHON %s --broken-numpy | FileCheck %s --check-prefix=BROKEN
+
+# Importing mlir.ir must not import NumPy. The dense elements attribute builders
+# import it when first called, and report a missing NumPy at that point.
+
+import sys
+
+import_error = None
+if "--without-numpy" in sys.argv:
+ sys.modules["numpy"] = None
+elif "--missing-dependency" in sys.argv:
+ import_error = ModuleNotFoundError(
+ "No module named 'numpy_dependency'", name="numpy_dependency"
+ )
+elif "--broken-numpy" in sys.argv:
+ import_error = ImportError("NumPy initialization failed")
+
+if import_error is not None:
+
+ class BrokenNumpyFinder:
+ def find_spec(self, fullname, path=None, target=None):
+ if fullname == "numpy":
+ raise import_error
+
+ sys.meta_path.insert(0, BrokenNumpyFinder())
+
+from mlir.ir import *
+
+# CHECK: numpy imported: False
+# NO-NUMPY: numpy imported: False
+# DEPENDENCY: numpy imported: False
+# BROKEN: numpy imported: False
+print("numpy imported:", sys.modules.get("numpy") is not None)
+
+# CHECK: F64ElementsAttr: dense<[1.500000e+00, 2.500000e+00]> : tensor<2xf64>
+# CHECK: I32ElementsAttr: dense<[1, 2, 3]> : tensor<3xi32>
+# CHECK: I64ElementsAttr: dense<[4, 5]> : tensor<2xi64>
+# CHECK: IndexElementsAttr: dense<6> : tensor<1xindex>
+# NO-NUMPY: F64ElementsAttr: ModuleNotFoundError: {{.*}}requires NumPy
+# NO-NUMPY: I32ElementsAttr: ModuleNotFoundError: {{.*}}requires NumPy
+# NO-NUMPY: I64ElementsAttr: ModuleNotFoundError: {{.*}}requires NumPy
+# NO-NUMPY: IndexElementsAttr: ModuleNotFoundError: {{.*}}requires NumPy
+# DEPENDENCY: F64ElementsAttr: ModuleNotFoundError: No module named 'numpy_dependency'
+# DEPENDENCY: I32ElementsAttr: ModuleNotFoundError: No module named 'numpy_dependency'
+# DEPENDENCY: I64ElementsAttr: ModuleNotFoundError: No module named 'numpy_dependency'
+# DEPENDENCY: IndexElementsAttr: ModuleNotFoundError: No module named 'numpy_dependency'
+# BROKEN: F64ElementsAttr: ImportError: NumPy initialization failed
+# BROKEN: I32ElementsAttr: ImportError: NumPy initialization failed
+# BROKEN: I64ElementsAttr: ImportError: NumPy initialization failed
+# BROKEN: IndexElementsAttr: ImportError: NumPy initialization failed
+with Context() as ctx:
+ for kind, values in (
+ ("F64ElementsAttr", [1.5, 2.5]),
+ ("I32ElementsAttr", [1, 2, 3]),
+ ("I64ElementsAttr", [4, 5]),
+ ("IndexElementsAttr", [6]),
+ ):
+ try:
+ print(f"{kind}:", AttrBuilder.get(kind)(values, context=ctx))
+ except ImportError as e:
+ if import_error is not None:
+ assert e is import_error
+ else:
+ assert "--without-numpy" in sys.argv
+ assert isinstance(e, ModuleNotFoundError) and e.name == "numpy"
+ print(f"{kind}: {type(e).__name__}: {e}")
+
+# CHECK: numpy imported: True
+# NO-NUMPY: numpy imported: False
+# DEPENDENCY: numpy imported: False
+# BROKEN: numpy imported: False
+print("numpy imported:", sys.modules.get("numpy") is not None)
``````````
</details>
https://github.com/llvm/llvm-project/pull/238232
More information about the Mlir-commits
mailing list