[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:04:41 PST 2026
https://github.com/brandon-yujie-sun created https://github.com/llvm/llvm-project/pull/238232
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
>From 202834c19d86d298e3076f14d2505d29468fa56e Mon Sep 17 00:00:00 2001
From: brandonsun <brandons at nvidia.com>
Date: Wed, 7 Oct 2026 02:15:05 -0700
Subject: [PATCH] [mlir][python] Import NumPy lazily in mlir.ir
Avoid importing NumPy during MLIR startup when no dense-elements attribute
builder is used. Register the four builders unconditionally and import NumPy
when called, with an actionable error if NumPy itself is unavailable. Preserve
missing transitive dependencies and other import failures unchanged.
Cover import laziness, all four builders, and dependency failure modes. Make
the memref test import NumPy explicitly instead of receiving np through a
wildcard import from mlir.ir.
---
mlir/python/mlir/ir.py | 93 +++++++++++++++++------------
mlir/test/python/dialects/memref.py | 2 +
mlir/test/python/ir/lazy_numpy.py | 75 +++++++++++++++++++++++
3 files changed, 133 insertions(+), 37 deletions(-)
create mode 100644 mlir/test/python/ir/lazy_numpy.py
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)
More information about the Mlir-commits
mailing list