[MLIR][Python] move Py* types into headers

mlir/lib/Bindings/Python/IRModule.h → mlir/include/mlir/Bindings/Python/IRModule.h

@@ -571,7 +571,7 @@ class PyOperationBase {
 
   /// Verify the operation. Throws `MLIRError` if verification fails, and
   /// returns `true` otherwise.
-  bool verify();
+  MLIR_CAPI_EXPORTED bool verify();
 
   /// Each must provide access to the raw Operation.
   virtual PyOperation &getOperation() = 0;

mlir/lib/Bindings/Python/DialectGPU.cpp

@@ -9,6 +9,7 @@
 #include "mlir-c/Dialect/GPU.h"
 #include "mlir-c/IR.h"
 #include "mlir-c/Support.h"
+#include "mlir/Bindings/Python/IRModule.h"
 #include "mlir/Bindings/Python/Nanobind.h"
 #include "mlir/Bindings/Python/NanobindAdaptors.h"
 
@@ -28,6 +29,7 @@ NB_MODULE(_mlirDialectsGPU, m) {
   //===-------------------------------------------------------------------===//
   // AsyncTokenType
   //===-------------------------------------------------------------------===//
+  m.def("blahblah", [](PyOperationBase &op) { op.verify(); });
 
   auto mlirGPUAsyncTokenType =
       mlir_type_subclass(m, "AsyncTokenType", mlirTypeIsAGPUAsyncTokenType);

mlir/lib/Bindings/Python/DialectSMT.cpp

@@ -6,15 +6,14 @@
 //
 //===----------------------------------------------------------------------===//
 
-#include "NanobindUtils.h"
-
 #include "mlir-c/Dialect/SMT.h"
 #include "mlir-c/IR.h"
 #include "mlir-c/Support.h"
 #include "mlir-c/Target/ExportSMTLIB.h"
 #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;
 

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(); },

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"
 

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;

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"

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;

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;

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;

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;

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 {

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;

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 {

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
@@ -889,3 +881,4 @@ add_mlir_python_modules(MLIRPythonModules
   COMMON_CAPI_LINK_LIBS
     MLIRPythonCAPI
 )
+add_dependencies(MLIRPythonModules.extension._mlirDialectsGPU.dso MLIRPythonModules.extension._mlir.dso)

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 *

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