[Mlir-commits] [mlir] [MLIR][Python] move Py* types into headers (PR #156575)
Maksim Levental
llvmlistbot at llvm.org
Mon Sep 8 20:53:03 PDT 2025
https://github.com/makslevental updated https://github.com/llvm/llvm-project/pull/156575
>From b779d77a6e7cd9c402c50c6be727be583cbf69fc Mon Sep 17 00:00:00 2001
From: makslevental <maksim.levental at gmail.com>
Date: Tue, 2 Sep 2025 20:03:22 -0700
Subject: [PATCH 1/2] [MLIR][Python] move Py* types
---
.../mlir}/Bindings/Python/Globals.h | 0
.../mlir}/Bindings/Python/IRModule.h | 0
.../mlir}/Bindings/Python/NanobindUtils.h | 0
mlir/lib/Bindings/Python/DialectSMT.cpp | 3 +-
mlir/lib/Bindings/Python/IRAffine.cpp | 48 ++++++++++---------
mlir/lib/Bindings/Python/IRAttributes.cpp | 4 +-
mlir/lib/Bindings/Python/IRCore.cpp | 13 +++--
mlir/lib/Bindings/Python/IRInterfaces.cpp | 2 +-
mlir/lib/Bindings/Python/IRModule.cpp | 12 +++--
mlir/lib/Bindings/Python/IRTypes.cpp | 5 +-
mlir/lib/Bindings/Python/MainModule.cpp | 6 +--
mlir/lib/Bindings/Python/Pass.cpp | 4 +-
mlir/lib/Bindings/Python/Pass.h | 2 +-
mlir/lib/Bindings/Python/Rewrite.cpp | 9 ++--
mlir/lib/Bindings/Python/Rewrite.h | 2 +-
mlir/python/CMakeLists.txt | 8 ----
16 files changed, 59 insertions(+), 59 deletions(-)
rename mlir/{lib => include/mlir}/Bindings/Python/Globals.h (100%)
rename mlir/{lib => include/mlir}/Bindings/Python/IRModule.h (100%)
rename mlir/{lib => include/mlir}/Bindings/Python/NanobindUtils.h (100%)
diff --git a/mlir/lib/Bindings/Python/Globals.h b/mlir/include/mlir/Bindings/Python/Globals.h
similarity index 100%
rename from mlir/lib/Bindings/Python/Globals.h
rename to mlir/include/mlir/Bindings/Python/Globals.h
diff --git a/mlir/lib/Bindings/Python/IRModule.h b/mlir/include/mlir/Bindings/Python/IRModule.h
similarity index 100%
rename from mlir/lib/Bindings/Python/IRModule.h
rename to mlir/include/mlir/Bindings/Python/IRModule.h
diff --git a/mlir/lib/Bindings/Python/NanobindUtils.h b/mlir/include/mlir/Bindings/Python/NanobindUtils.h
similarity index 100%
rename from mlir/lib/Bindings/Python/NanobindUtils.h
rename to mlir/include/mlir/Bindings/Python/NanobindUtils.h
diff --git a/mlir/lib/Bindings/Python/DialectSMT.cpp b/mlir/lib/Bindings/Python/DialectSMT.cpp
index 3123e3bdda496..7cb8776b78ed2 100644
--- a/mlir/lib/Bindings/Python/DialectSMT.cpp
+++ b/mlir/lib/Bindings/Python/DialectSMT.cpp
@@ -6,8 +6,6 @@
//
//===----------------------------------------------------------------------===//
-#include "NanobindUtils.h"
-
#include "mlir-c/Dialect/SMT.h"
#include "mlir-c/IR.h"
#include "mlir-c/Support.h"
@@ -15,6 +13,7 @@
#include "mlir/Bindings/Python/Diagnostics.h"
#include "mlir/Bindings/Python/Nanobind.h"
#include "mlir/Bindings/Python/NanobindAdaptors.h"
+#include "mlir/Bindings/Python/NanobindUtils.h"
namespace nb = nanobind;
diff --git a/mlir/lib/Bindings/Python/IRAffine.cpp b/mlir/lib/Bindings/Python/IRAffine.cpp
index bc6aa0dac6221..faa9b489c8e90 100644
--- a/mlir/lib/Bindings/Python/IRAffine.cpp
+++ b/mlir/lib/Bindings/Python/IRAffine.cpp
@@ -13,19 +13,22 @@
#include <utility>
#include <vector>
-#include "IRModule.h"
-#include "NanobindUtils.h"
#include "mlir-c/AffineExpr.h"
#include "mlir-c/AffineMap.h"
-#include "mlir-c/Bindings/Python/Interop.h" // This is expected after nanobind.
#include "mlir-c/IntegerSet.h"
-#include "mlir/Bindings/Python/Nanobind.h"
+#include "mlir/Bindings/Python/IRModule.h"
+#include "mlir/Bindings/Python/NanobindUtils.h"
#include "mlir/Support/LLVM.h"
#include "llvm/ADT/Hashing.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/ADT/StringRef.h"
#include "llvm/ADT/Twine.h"
+// clang-format off
+#include "mlir/Bindings/Python/Nanobind.h"
+#include "mlir-c/Bindings/Python/Interop.h" // ON WINDOWS This is expected after nanobind.
+// clang-format on
+
namespace nb = nanobind;
using namespace mlir;
using namespace mlir::python;
@@ -706,25 +709,24 @@ void mlir::python::populateIRAffine(nb::module_ &m) {
[](PyAffineMap &self) {
return static_cast<size_t>(llvm::hash_value(self.get().ptr));
})
- .def_static("compress_unused_symbols",
- [](const nb::list &affineMaps,
- DefaultingPyMlirContext context) {
- SmallVector<MlirAffineMap> maps;
- pyListToVector<PyAffineMap, MlirAffineMap>(
- affineMaps, maps, "attempting to create an AffineMap");
- std::vector<MlirAffineMap> compressed(affineMaps.size());
- auto populate = [](void *result, intptr_t idx,
- MlirAffineMap m) {
- static_cast<MlirAffineMap *>(result)[idx] = (m);
- };
- mlirAffineMapCompressUnusedSymbols(
- maps.data(), maps.size(), compressed.data(), populate);
- std::vector<PyAffineMap> res;
- res.reserve(compressed.size());
- for (auto m : compressed)
- res.emplace_back(context->getRef(), m);
- return res;
- })
+ .def_static(
+ "compress_unused_symbols",
+ [](const nb::list &affineMaps, DefaultingPyMlirContext context) {
+ SmallVector<MlirAffineMap> maps;
+ pyListToVector<PyAffineMap, MlirAffineMap>(
+ affineMaps, maps, "attempting to create an AffineMap");
+ std::vector<MlirAffineMap> compressed(affineMaps.size());
+ auto populate = [](void *result, intptr_t idx, MlirAffineMap m) {
+ static_cast<MlirAffineMap *>(result)[idx] = (m);
+ };
+ mlirAffineMapCompressUnusedSymbols(maps.data(), maps.size(),
+ compressed.data(), populate);
+ std::vector<PyAffineMap> res;
+ res.reserve(compressed.size());
+ for (auto m : compressed)
+ res.emplace_back(context->getRef(), m);
+ return res;
+ })
.def_prop_ro(
"context",
[](PyAffineMap &self) { return self.getContext().getObject(); },
diff --git a/mlir/lib/Bindings/Python/IRAttributes.cpp b/mlir/lib/Bindings/Python/IRAttributes.cpp
index b3c768846c74f..da425b0c0e218 100644
--- a/mlir/lib/Bindings/Python/IRAttributes.cpp
+++ b/mlir/lib/Bindings/Python/IRAttributes.cpp
@@ -12,12 +12,12 @@
#include <string_view>
#include <utility>
-#include "IRModule.h"
-#include "NanobindUtils.h"
#include "mlir-c/BuiltinAttributes.h"
#include "mlir-c/BuiltinTypes.h"
+#include "mlir/Bindings/Python/IRModule.h"
#include "mlir/Bindings/Python/Nanobind.h"
#include "mlir/Bindings/Python/NanobindAdaptors.h"
+#include "mlir/Bindings/Python/NanobindUtils.h"
#include "llvm/ADT/ScopeExit.h"
#include "llvm/Support/raw_ostream.h"
diff --git a/mlir/lib/Bindings/Python/IRCore.cpp b/mlir/lib/Bindings/Python/IRCore.cpp
index 8273a9346e5dd..a09d7658feff4 100644
--- a/mlir/lib/Bindings/Python/IRCore.cpp
+++ b/mlir/lib/Bindings/Python/IRCore.cpp
@@ -6,21 +6,24 @@
//
//===----------------------------------------------------------------------===//
-#include "Globals.h"
-#include "IRModule.h"
-#include "NanobindUtils.h"
-#include "mlir-c/Bindings/Python/Interop.h" // This is expected after nanobind.
#include "mlir-c/BuiltinAttributes.h"
#include "mlir-c/Debug.h"
#include "mlir-c/Diagnostics.h"
#include "mlir-c/IR.h"
#include "mlir-c/Support.h"
-#include "mlir/Bindings/Python/Nanobind.h"
+#include "mlir/Bindings/Python/Globals.h"
+#include "mlir/Bindings/Python/IRModule.h"
#include "mlir/Bindings/Python/NanobindAdaptors.h"
+#include "mlir/Bindings/Python/NanobindUtils.h"
#include "nanobind/nanobind.h"
#include "llvm/ADT/ArrayRef.h"
#include "llvm/ADT/SmallVector.h"
+// clang-format off
+#include "mlir/Bindings/Python/Nanobind.h"
+#include "mlir-c/Bindings/Python/Interop.h" // ON WINDOWS This is expected after nanobind.
+// clang-format on
+
#include <optional>
namespace nb = nanobind;
diff --git a/mlir/lib/Bindings/Python/IRInterfaces.cpp b/mlir/lib/Bindings/Python/IRInterfaces.cpp
index 6aa057c2a78aa..89d42909a0865 100644
--- a/mlir/lib/Bindings/Python/IRInterfaces.cpp
+++ b/mlir/lib/Bindings/Python/IRInterfaces.cpp
@@ -12,11 +12,11 @@
#include <utility>
#include <vector>
-#include "IRModule.h"
#include "mlir-c/BuiltinAttributes.h"
#include "mlir-c/IR.h"
#include "mlir-c/Interfaces.h"
#include "mlir-c/Support.h"
+#include "mlir/Bindings/Python/IRModule.h"
#include "mlir/Bindings/Python/Nanobind.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/SmallVector.h"
diff --git a/mlir/lib/Bindings/Python/IRModule.cpp b/mlir/lib/Bindings/Python/IRModule.cpp
index 0de2f1711829b..28bb73832f79e 100644
--- a/mlir/lib/Bindings/Python/IRModule.cpp
+++ b/mlir/lib/Bindings/Python/IRModule.cpp
@@ -6,16 +6,18 @@
//
//===----------------------------------------------------------------------===//
-#include "IRModule.h"
-
#include <optional>
#include <vector>
-#include "Globals.h"
-#include "NanobindUtils.h"
-#include "mlir-c/Bindings/Python/Interop.h"
#include "mlir-c/Support.h"
+#include "mlir/Bindings/Python/Globals.h"
+#include "mlir/Bindings/Python/IRModule.h"
+#include "mlir/Bindings/Python/NanobindUtils.h"
+
+// clang-format off
#include "mlir/Bindings/Python/Nanobind.h"
+#include "mlir-c/Bindings/Python/Interop.h" // ON WINDOWS This is expected after nanobind.
+// clang-format on
namespace nb = nanobind;
using namespace mlir;
diff --git a/mlir/lib/Bindings/Python/IRTypes.cpp b/mlir/lib/Bindings/Python/IRTypes.cpp
index 745b9c48c5165..3730eeb8b40f3 100644
--- a/mlir/lib/Bindings/Python/IRTypes.cpp
+++ b/mlir/lib/Bindings/Python/IRTypes.cpp
@@ -7,17 +7,16 @@
//===----------------------------------------------------------------------===//
// clang-format off
-#include "IRModule.h"
+#include "mlir/Bindings/Python/IRModule.h"
#include "mlir/Bindings/Python/IRTypes.h"
// clang-format on
#include <optional>
-#include "IRModule.h"
-#include "NanobindUtils.h"
#include "mlir-c/BuiltinAttributes.h"
#include "mlir-c/BuiltinTypes.h"
#include "mlir-c/Support.h"
+#include "mlir/Bindings/Python/NanobindUtils.h"
namespace nb = nanobind;
using namespace mlir;
diff --git a/mlir/lib/Bindings/Python/MainModule.cpp b/mlir/lib/Bindings/Python/MainModule.cpp
index d7282b3d6f713..f9ae11039247e 100644
--- a/mlir/lib/Bindings/Python/MainModule.cpp
+++ b/mlir/lib/Bindings/Python/MainModule.cpp
@@ -6,12 +6,12 @@
//
//===----------------------------------------------------------------------===//
-#include "Globals.h"
-#include "IRModule.h"
-#include "NanobindUtils.h"
#include "Pass.h"
#include "Rewrite.h"
+#include "mlir/Bindings/Python/Globals.h"
+#include "mlir/Bindings/Python/IRModule.h"
#include "mlir/Bindings/Python/Nanobind.h"
+#include "mlir/Bindings/Python/NanobindUtils.h"
namespace nb = nanobind;
using namespace mlir;
diff --git a/mlir/lib/Bindings/Python/Pass.cpp b/mlir/lib/Bindings/Python/Pass.cpp
index 6ee85e8a31492..dbf766d3f6356 100644
--- a/mlir/lib/Bindings/Python/Pass.cpp
+++ b/mlir/lib/Bindings/Python/Pass.cpp
@@ -8,11 +8,11 @@
#include "Pass.h"
-#include "IRModule.h"
#include "mlir-c/Pass.h"
+#include "mlir/Bindings/Python/IRModule.h"
// clang-format off
#include "mlir/Bindings/Python/Nanobind.h"
-#include "mlir-c/Bindings/Python/Interop.h" // This is expected after nanobind.
+#include "mlir-c/Bindings/Python/Interop.h" // ON WINDOWS This is expected after nanobind.
// clang-format on
namespace nb = nanobind;
diff --git a/mlir/lib/Bindings/Python/Pass.h b/mlir/lib/Bindings/Python/Pass.h
index bc40943521829..0221bd10e723e 100644
--- a/mlir/lib/Bindings/Python/Pass.h
+++ b/mlir/lib/Bindings/Python/Pass.h
@@ -9,7 +9,7 @@
#ifndef MLIR_BINDINGS_PYTHON_PASS_H
#define MLIR_BINDINGS_PYTHON_PASS_H
-#include "NanobindUtils.h"
+#include "mlir/Bindings/Python/NanobindUtils.h"
namespace mlir {
namespace python {
diff --git a/mlir/lib/Bindings/Python/Rewrite.cpp b/mlir/lib/Bindings/Python/Rewrite.cpp
index 5b7de50f02e6a..74453b45ebace 100644
--- a/mlir/lib/Bindings/Python/Rewrite.cpp
+++ b/mlir/lib/Bindings/Python/Rewrite.cpp
@@ -8,12 +8,15 @@
#include "Rewrite.h"
-#include "IRModule.h"
#include "mlir-c/Rewrite.h"
-#include "mlir/Bindings/Python/Nanobind.h"
-#include "mlir-c/Bindings/Python/Interop.h" // This is expected after nanobind.
+#include "mlir/Bindings/Python/IRModule.h"
#include "mlir/Config/mlir-config.h"
+// clang-format off
+#include "mlir/Bindings/Python/Nanobind.h"
+#include "mlir-c/Bindings/Python/Interop.h" // ON WINDOWS This is expected after nanobind.
+// clang-format on
+
namespace nb = nanobind;
using namespace mlir;
using namespace nb::literals;
diff --git a/mlir/lib/Bindings/Python/Rewrite.h b/mlir/lib/Bindings/Python/Rewrite.h
index ae89e2b9589f1..f8ffdc7bdc458 100644
--- a/mlir/lib/Bindings/Python/Rewrite.h
+++ b/mlir/lib/Bindings/Python/Rewrite.h
@@ -9,7 +9,7 @@
#ifndef MLIR_BINDINGS_PYTHON_REWRITE_H
#define MLIR_BINDINGS_PYTHON_REWRITE_H
-#include "NanobindUtils.h"
+#include "mlir/Bindings/Python/NanobindUtils.h"
namespace mlir {
namespace python {
diff --git a/mlir/python/CMakeLists.txt b/mlir/python/CMakeLists.txt
index 8e7949480f21e..6d61d891fc65a 100644
--- a/mlir/python/CMakeLists.txt
+++ b/mlir/python/CMakeLists.txt
@@ -490,12 +490,6 @@ declare_mlir_python_extension(MLIRPythonExtension.Core
Pass.cpp
Rewrite.cpp
- # Headers must be included explicitly so they are installed.
- Globals.h
- IRModule.h
- Pass.h
- NanobindUtils.h
- Rewrite.h
PRIVATE_LINK_LIBS
LLVMSupport
EMBED_CAPI_LINK_LIBS
@@ -732,8 +726,6 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.SMT.Pybind
PYTHON_BINDINGS_LIBRARY nanobind
SOURCES
DialectSMT.cpp
- # Headers must be included explicitly so they are installed.
- NanobindUtils.h
PRIVATE_LINK_LIBS
LLVMSupport
EMBED_CAPI_LINK_LIBS
>From 0da6975869270e9c7daf9e155682a85a0b0ffdea Mon Sep 17 00:00:00 2001
From: Maksim Levental <maksim.levental at gmail.com>
Date: Mon, 8 Sep 2025 22:52:27 -0400
Subject: [PATCH 2/2] test PyOperationBase method call
---
mlir/lib/Bindings/Python/DialectGPU.cpp | 4 ++++
mlir/python/CMakeLists.txt | 2 ++
mlir/python/mlir/dialects/gpu/__init__.py | 1 +
mlir/test/python/dialects/gpu/dialect.py | 3 ++-
4 files changed, 9 insertions(+), 1 deletion(-)
diff --git a/mlir/lib/Bindings/Python/DialectGPU.cpp b/mlir/lib/Bindings/Python/DialectGPU.cpp
index 2568d535edb5a..60d832978de96 100644
--- a/mlir/lib/Bindings/Python/DialectGPU.cpp
+++ b/mlir/lib/Bindings/Python/DialectGPU.cpp
@@ -11,6 +11,7 @@
#include "mlir-c/Support.h"
#include "mlir/Bindings/Python/Nanobind.h"
#include "mlir/Bindings/Python/NanobindAdaptors.h"
+#include "mlir/Bindings/Python/IRModule.h"
namespace nb = nanobind;
using namespace nanobind::literals;
@@ -28,6 +29,9 @@ NB_MODULE(_mlirDialectsGPU, m) {
//===-------------------------------------------------------------------===//
// AsyncTokenType
//===-------------------------------------------------------------------===//
+ m.def("blahblah", [](PyOperationBase& op) {
+ op.verify();
+ });
auto mlirGPUAsyncTokenType =
mlir_type_subclass(m, "AsyncTokenType", mlirTypeIsAGPUAsyncTokenType);
diff --git a/mlir/python/CMakeLists.txt b/mlir/python/CMakeLists.txt
index 6d61d891fc65a..50f581f0e8b6a 100644
--- a/mlir/python/CMakeLists.txt
+++ b/mlir/python/CMakeLists.txt
@@ -881,3 +881,5 @@ add_mlir_python_modules(MLIRPythonModules
COMMON_CAPI_LINK_LIBS
MLIRPythonCAPI
)
+set_property(TARGET MLIRPythonModules.extension._mlir.dso PROPERTY CXX_VISIBILITY_PRESET default)
+add_dependencies(MLIRPythonModules.extension._mlirDialectsGPU.dso MLIRPythonModules.extension._mlir.dso)
diff --git a/mlir/python/mlir/dialects/gpu/__init__.py b/mlir/python/mlir/dialects/gpu/__init__.py
index 4cd80aa8b7ca8..460957cc37569 100644
--- a/mlir/python/mlir/dialects/gpu/__init__.py
+++ b/mlir/python/mlir/dialects/gpu/__init__.py
@@ -2,6 +2,7 @@
# See https://llvm.org/LICENSE.txt for license information.
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+from ..._mlir_libs import _mlir
from .._gpu_ops_gen import *
from .._gpu_enum_gen import *
from ..._mlir_libs._mlirDialectsGPU import *
diff --git a/mlir/test/python/dialects/gpu/dialect.py b/mlir/test/python/dialects/gpu/dialect.py
index 26ee9f34cb332..6454b9c999781 100644
--- a/mlir/test/python/dialects/gpu/dialect.py
+++ b/mlir/test/python/dialects/gpu/dialect.py
@@ -26,7 +26,8 @@ def testGPUPass():
def testMMAElementWiseAttr():
module = Module.create()
with InsertionPoint(module.body):
- gpu.BlockDimOp(gpu.Dimension.y)
+ b = gpu.BlockDimOp(gpu.Dimension.y)
+ gpu.blahblah(b)
# CHECK: %block_dim_y = gpu.block_dim y
print(module)
pass
More information about the Mlir-commits
mailing list