[mlir][python] Apply py::module_local() to all classes.

[stellaraccident]

Summary

• This allows multiple MLIR-API embedding downstreams to co-exist in the same process.
• I believe this is the last thing needed to enable isolated embedding.

mlir/lib/Bindings/Python/ExecutionEngineModule.cpp

 
 /// Create the `mlir.execution_engine` module here.
 PYBIND11_MODULE(_mlirExecutionEngine, m) {
   m.doc() = "MLIR Execution Engine";
 
   //----------------------------------------------------------------------------
   // Mapping of the top-level PassManager
   //----------------------------------------------------------------------------
-  py::class_<PyExecutionEngine>(m, "ExecutionEngine")
+  py::class_<PyExecutionEngine>(m, "ExecutionEngine", py::module_local())
       .def(py::init<>([](MlirModule module, int optLevel,
                          const std::vector<std::string> &sharedLibPaths) {
              llvm::SmallVector<MlirStringRef, 4> libPaths;
              for (const std::string &path : sharedLibPaths)
                libPaths.push_back({path.c_str(), path.length()});
              MlirExecutionEngine executionEngine = mlirExecutionEngineCreate(
                  module, optLevel, libPaths.size(), libPaths.data());
              if (mlirExecutionEngineIsNull(executionEngine))

mlir/lib/Bindings/Python/IRAffine.cpp

if (!DerivedTy::isaFunction(orig)) {
       throw SetPyError(PyExc_ValueError,
                        Twine("Cannot cast affine expression to ") +
                            DerivedTy::pyClassName + " (from " + origRepr + ")");
     }
     return orig;
   }
 
   static void bind(py::module &m) {
-    auto cls = ClassTy(m, DerivedTy::pyClassName);
+    auto cls = ClassTy(m, DerivedTy::pyClassName, py::module_local());
     cls.def(py::init<PyAffineExpr &>());
     DerivedTy::bindDerived(cls);
   }
 
   /// Implemented by derived classes to add methods to the Python subclass.
   static void bindDerived(ClassTy &m) {}
 };
 
public:
   PyAffineExpr getExpr() {
     return PyAffineExpr(set.getContext(),
                         mlirIntegerSetGetConstraint(set, pos));
   }
 
   bool isEq() { return mlirIntegerSetIsConstraintEq(set, pos); }
 
   static void bind(py::module &m) {
-    py::class_<PyIntegerSetConstraint>(m, "IntegerSetConstraint")
+    py::class_<PyIntegerSetConstraint>(m, "IntegerSetConstraint",
+                                       py::module_local())
         .def_property_readonly("expr", &PyIntegerSetConstraint::getExpr)
         .def_property_readonly("is_eq", &PyIntegerSetConstraint::isEq);
   }
 
 private:
   PyIntegerSet set;
   intptr_t pos;
 };
return PyIntegerSet(
       PyMlirContext::forContext(mlirIntegerSetGetContext(rawIntegerSet)),
       rawIntegerSet);
 }
 
 void mlir::python::populateIRAffine(py::module &m) {
   //----------------------------------------------------------------------------
   // Mapping of PyAffineExpr and derived classes.
   //----------------------------------------------------------------------------
-  py::class_<PyAffineExpr>(m, "AffineExpr")
+  py::class_<PyAffineExpr>(m, "AffineExpr", py::module_local())
       .def_property_readonly(MLIR_PYTHON_CAPI_PTR_ATTR,
                              &PyAffineExpr::getCapsule)
       .def(MLIR_PYTHON_CAPI_FACTORY_ATTR, &PyAffineExpr::createFromCapsule)
       .def("__add__",
            [](PyAffineExpr &self, PyAffineExpr &other) {
              return PyAffineAddExpr::get(self, other);
            })
       .def("__mul__",
void mlir::python::populateIRAffine(py::module &m) {
   PyAffineMulExpr::bind(m);
   PyAffineModExpr::bind(m);
   PyAffineFloorDivExpr::bind(m);
   PyAffineCeilDivExpr::bind(m);
 
   //----------------------------------------------------------------------------
   // Mapping of PyAffineMap.
   //----------------------------------------------------------------------------
-  py::class_<PyAffineMap>(m, "AffineMap")
+  py::class_<PyAffineMap>(m, "AffineMap", py::module_local())
       .def_property_readonly(MLIR_PYTHON_CAPI_PTR_ATTR,
                              &PyAffineMap::getCapsule)
       .def(MLIR_PYTHON_CAPI_FACTORY_ATTR, &PyAffineMap::createFromCapsule)
       .def("__eq__",
            [](PyAffineMap &self, PyAffineMap &other) { return self == other; })
       .def("__eq__", [](PyAffineMap &self, py::object &other) { return false; })
       .def("__str__",
            [](PyAffineMap &self) {
py::class_<PyAffineMap>(m, "AffineMap", py::module_local())
       .def_property_readonly("results", [](PyAffineMap &self) {
         return PyAffineMapExprList(self);
       });
   PyAffineMapExprList::bind(m);
 
   //----------------------------------------------------------------------------
   // Mapping of PyIntegerSet.
   //----------------------------------------------------------------------------
-  py::class_<PyIntegerSet>(m, "IntegerSet")
+  py::class_<PyIntegerSet>(m, "IntegerSet", py::module_local())
       .def_property_readonly(MLIR_PYTHON_CAPI_PTR_ATTR,
                              &PyIntegerSet::getCapsule)
       .def(MLIR_PYTHON_CAPI_FACTORY_ATTR, &PyIntegerSet::createFromCapsule)
       .def("__eq__", [](PyIntegerSet &self,
                         PyIntegerSet &other) { return self == other; })
       .def("__eq__", [](PyIntegerSet &self, py::object other) { return false; })
       .def("__str__",
            [](PyIntegerSet &self) {

mlir/lib/Bindings/Python/IRAttributes.cpp

PyAttribute dunderNext() {
       if (nextIndex >= mlirArrayAttrGetNumElements(attr.get())) {
         throw py::stop_iteration();
       }
       return PyAttribute(attr.getContext(),
                          mlirArrayAttrGetElement(attr.get(), nextIndex++));
     }
 
     static void bind(py::module &m) {
-      py::class_<PyArrayAttributeIterator>(m, "ArrayAttributeIterator")
+      py::class_<PyArrayAttributeIterator>(m, "ArrayAttributeIterator",
+                                           py::module_local())
           .def("__iter__", &PyArrayAttributeIterator::dunderIter)
           .def("__next__", &PyArrayAttributeIterator::dunderNext);
     }
 
   private:
     PyAttribute attr;
     int nextIndex = 0;
   };

mlir/lib/Bindings/Python/IRCore.cpp

 /// Wrapper for the global LLVM debugging flag.
 struct PyGlobalDebugFlag {
   static void set(py::object &o, bool enable) { mlirEnableGlobalDebug(enable); }
 
   static bool get(py::object) { return mlirIsGlobalDebugEnabled(); }
 
   static void bind(py::module &m) {
     // Debug flags.
-    py::class_<PyGlobalDebugFlag>(m, "_GlobalDebug")
+    py::class_<PyGlobalDebugFlag>(m, "_GlobalDebug", py::module_local())
         .def_property_static("flag", &PyGlobalDebugFlag::get,
                              &PyGlobalDebugFlag::set, "LLVM-wide debug flag");
   }
 };
 
 //------------------------------------------------------------------------------
 // Collections.
 //------------------------------------------------------------------------------
PyRegion dunderNext() {
     if (nextIndex >= mlirOperationGetNumRegions(operation->get())) {
       throw py::stop_iteration();
     }
     MlirRegion region = mlirOperationGetRegion(operation->get(), nextIndex++);
     return PyRegion(operation, region);
   }
 
   static void bind(py::module &m) {
-    py::class_<PyRegionIterator>(m, "RegionIterator")
+    py::class_<PyRegionIterator>(m, "RegionIterator", py::module_local())
         .def("__iter__", &PyRegionIterator::dunderIter)
         .def("__next__", &PyRegionIterator::dunderNext);
   }
 
 private:
   PyOperationRef operation;
   int nextIndex = 0;
 };
if (index < 0 || index >= dunderLen()) {
       throw SetPyError(PyExc_IndexError,
                        "attempt to access out of bounds region");
     }
     MlirRegion region = mlirOperationGetRegion(operation->get(), index);
     return PyRegion(operation, region);
   }
 
   static void bind(py::module &m) {
-    py::class_<PyRegionList>(m, "RegionSequence")
+    py::class_<PyRegionList>(m, "RegionSequence", py::module_local())
         .def("__len__", &PyRegionList::dunderLen)
         .def("__getitem__", &PyRegionList::dunderGetItem);
   }
 
 private:
   PyOperationRef operation;
 };
 
PyBlock dunderNext() {
     }
 
     PyBlock returnBlock(operation, next);
     next = mlirBlockGetNextInRegion(next);
     return returnBlock;
   }
 
   static void bind(py::module &m) {
-    py::class_<PyBlockIterator>(m, "BlockIterator")
+    py::class_<PyBlockIterator>(m, "BlockIterator", py::module_local())
         .def("__iter__", &PyBlockIterator::dunderIter)
         .def("__next__", &PyBlockIterator::dunderNext);
   }
 
 private:
   PyOperationRef operation;
   MlirBlock next;
 };
PyBlock appendBlock(py::args pyArgTypes) {
     }
 
     MlirBlock block = mlirBlockCreate(argTypes.size(), argTypes.data());
     mlirRegionAppendOwnedBlock(region, block);
     return PyBlock(operation, block);
   }
 
   static void bind(py::module &m) {
-    py::class_<PyBlockList>(m, "BlockList")
+    py::class_<PyBlockList>(m, "BlockList", py::module_local())
         .def("__getitem__", &PyBlockList::dunderGetItem)
         .def("__iter__", &PyBlockList::dunderIter)
         .def("__len__", &PyBlockList::dunderLen)
         .def("append", &PyBlockList::appendBlock, kAppendBlockDocstring);
   }
 
 private:
   PyOperationRef operation;
py::object dunderNext() {
 
     PyOperationRef returnOperation =
         PyOperation::forOperation(parentOperation->getContext(), next);
     next = mlirOperationGetNextInBlock(next);
     return returnOperation->createOpView();
   }
 
   static void bind(py::module &m) {
-    py::class_<PyOperationIterator>(m, "OperationIterator")
+    py::class_<PyOperationIterator>(m, "OperationIterator", py::module_local())
         .def("__iter__", &PyOperationIterator::dunderIter)
         .def("__next__", &PyOperationIterator::dunderNext);
   }
 
 private:
   PyOperationRef parentOperation;
   MlirOperation next;
 };
while (!mlirOperationIsNull(childOp)) {
       childOp = mlirOperationGetNextInBlock(childOp);
       index -= 1;
     }
     throw SetPyError(PyExc_IndexError,
                      "attempt to access out of bounds operation");
   }
 
   static void bind(py::module &m) {
-    py::class_<PyOperationList>(m, "OperationList")
+    py::class_<PyOperationList>(m, "OperationList", py::module_local())
         .def("__getitem__", &PyOperationList::dunderGetItem)
         .def("__iter__", &PyOperationList::dunderIter)
         .def("__len__", &PyOperationList::dunderLen);
   }
 
 private:
   PyOperationRef parentOperation;
   MlirBlock block;
if (!DerivedTy::isaFunction(orig.get())) {
                                              DerivedTy::pyClassName +
                                              " (from " + origRepr + ")");
     }
     return orig.get();
   }
 
   /// Binds the Python module objects to functions of this class.
   static void bind(py::module &m) {
-    auto cls = ClassTy(m, DerivedTy::pyClassName);
+    auto cls = ClassTy(m, DerivedTy::pyClassName, py::module_local());
     cls.def(py::init<PyValue &>(), py::keep_alive<0, 1>());
     DerivedTy::bindDerived(cls);
   }
 
   /// Implemented by derived classes to add methods to the Python subclass.
   static void bindDerived(ClassTy &m) {}
 };
 
if (index < 0 || index >= dunderLen()) {
                        "attempt to access out of bounds region");
     }
     PyValue value(operation, mlirBlockGetArgument(block, index));
     return PyBlockArgument(value);
   }
 
   /// Defines a Python class in the bindings.
   static void bind(py::module &m) {
-    py::class_<PyBlockArgumentList>(m, "BlockArgumentList")
+    py::class_<PyBlockArgumentList>(m, "BlockArgumentList", py::module_local())
         .def("__len__", &PyBlockArgumentList::dunderLen)
         .def("__getitem__", &PyBlockArgumentList::dunderGetItem);
   }
 
 private:
   PyOperationRef operation;
   MlirBlock block;
 };
public:
   }
 
   bool dunderContains(const std::string &name) {
     return !mlirAttributeIsNull(mlirOperationGetAttributeByName(
         operation->get(), toMlirStringRef(name)));
   }
 
   static void bind(py::module &m) {
-    py::class_<PyOpAttributeMap>(m, "OpAttributeMap")
+    py::class_<PyOpAttributeMap>(m, "OpAttributeMap", py::module_local())
         .def("__contains__", &PyOpAttributeMap::dunderContains)
         .def("__len__", &PyOpAttributeMap::dunderLen)
         .def("__getitem__", &PyOpAttributeMap::dunderGetItemNamed)
         .def("__getitem__", &PyOpAttributeMap::dunderGetItemIndexed)
         .def("__setitem__", &PyOpAttributeMap::dunderSetItem)
         .def("__delitem__", &PyOpAttributeMap::dunderDelItem);
   }
 
 private:
   PyOperationRef operation;
 };
 
 } // end namespace
 
 //------------------------------------------------------------------------------
 // Populates the core exports of the 'ir' submodule.
 //------------------------------------------------------------------------------
 
 void mlir::python::populateIRCore(py::module &m) {
   //----------------------------------------------------------------------------
   // Mapping of MlirContext.
   //----------------------------------------------------------------------------
-  py::class_<PyMlirContext>(m, "Context")
+  py::class_<PyMlirContext>(m, "Context", py::module_local())
       .def(py::init<>(&PyMlirContext::createNewContextForInit))
       .def_static("_get_live_count", &PyMlirContext::getLiveCount)
       .def("_get_context_again",
            [](PyMlirContext &self) {
              PyMlirContextRef ref = PyMlirContext::forContext(self.get());
              return ref.releaseObject();
            })
       .def("_get_live_operation_count", &PyMlirContext::getLiveOperationCount)
py::class_<PyMlirContext>(m, "Context", py::module_local())
            [](PyMlirContext &self, std::string &name) {
              return mlirContextIsRegisteredOperation(
                  self.get(), MlirStringRef{name.data(), name.size()});
            });
 
   //----------------------------------------------------------------------------
   // Mapping of PyDialectDescriptor
   //----------------------------------------------------------------------------
-  py::class_<PyDialectDescriptor>(m, "DialectDescriptor")
+  py::class_<PyDialectDescriptor>(m, "DialectDescriptor", py::module_local())
       .def_property_readonly("namespace",
                              [](PyDialectDescriptor &self) {
                                MlirStringRef ns =
                                    mlirDialectGetNamespace(self.get());
                                return py::str(ns.data, ns.length);
                              })
       .def("__repr__", [](PyDialectDescriptor &self) {
         MlirStringRef ns = mlirDialectGetNamespace(self.get());
         std::string repr("<DialectDescriptor ");
         repr.append(ns.data, ns.length);
         repr.append(">");
         return repr;
       });
 
   //----------------------------------------------------------------------------
   // Mapping of PyDialects
   //----------------------------------------------------------------------------
-  py::class_<PyDialects>(m, "Dialects")
+  py::class_<PyDialects>(m, "Dialects", py::module_local())
       .def("__getitem__",
            [=](PyDialects &self, std::string keyName) {
              MlirDialect dialect =
                  self.getDialectForKey(keyName, /*attrError=*/false);
              py::object descriptor =
                  py::cast(PyDialectDescriptor{self.getContext(), dialect});
              return createCustomDialectWrapper(keyName, std::move(descriptor));
            })
       .def("__getattr__", [=](PyDialects &self, std::string attrName) {
         MlirDialect dialect =
             self.getDialectForKey(attrName, /*attrError=*/true);
         py::object descriptor =
             py::cast(PyDialectDescriptor{self.getContext(), dialect});
         return createCustomDialectWrapper(attrName, std::move(descriptor));
       });
 
   //----------------------------------------------------------------------------
   // Mapping of PyDialect
   //----------------------------------------------------------------------------
-  py::class_<PyDialect>(m, "Dialect")
+  py::class_<PyDialect>(m, "Dialect", py::module_local())
       .def(py::init<py::object>(), "descriptor")
       .def_property_readonly(
           "descriptor", [](PyDialect &self) { return self.getDescriptor(); })
       .def("__repr__", [](py::object self) {
         auto clazz = self.attr("__class__");
         return py::str("<Dialect ") +
                self.attr("descriptor").attr("namespace") + py::str(" (class ") +
                clazz.attr("__module__") + py::str(".") +
                clazz.attr("__name__") + py::str(")>");
       });
 
   //----------------------------------------------------------------------------
   // Mapping of Location
   //----------------------------------------------------------------------------
-  py::class_<PyLocation>(m, "Location")
+  py::class_<PyLocation>(m, "Location", py::module_local())
       .def_property_readonly(MLIR_PYTHON_CAPI_PTR_ATTR, &PyLocation::getCapsule)
       .def(MLIR_PYTHON_CAPI_FACTORY_ATTR, &PyLocation::createFromCapsule)
       .def("__enter__", &PyLocation::contextEnter)
       .def("__exit__", &PyLocation::contextExit)
       .def("__eq__",
            [](PyLocation &self, PyLocation &other) -> bool {
              return mlirLocationEqual(self, other);
            })
py::class_<PyLocation>(m, "Location", py::module_local())
         mlirLocationPrint(self, printAccum.getCallback(),
                           printAccum.getUserData());
         return printAccum.join();
       });
 
   //----------------------------------------------------------------------------
   // Mapping of Module
   //----------------------------------------------------------------------------
-  py::class_<PyModule>(m, "Module")
+  py::class_<PyModule>(m, "Module", py::module_local())
       .def_property_readonly(MLIR_PYTHON_CAPI_PTR_ATTR, &PyModule::getCapsule)
       .def(MLIR_PYTHON_CAPI_FACTORY_ATTR, &PyModule::createFromCapsule)
       .def_static(
           "parse",
           [](const std::string moduleAsm, DefaultingPyMlirContext context) {
             MlirModule module = mlirModuleCreateParse(
                 context->get(), toMlirStringRef(moduleAsm));
             // TODO: Rework error reporting once diagnostic engine is exposed
py::class_<PyModule>(m, "Module", py::module_local())
                                printAccum.getUserData());
             return printAccum.join();
           },
           kOperationStrDunderDocstring);
 
   //----------------------------------------------------------------------------
   // Mapping of Operation.
   //----------------------------------------------------------------------------
-  py::class_<PyOperationBase>(m, "_OperationBase")
+  py::class_<PyOperationBase>(m, "_OperationBase", py::module_local())
       .def("__eq__",
            [](PyOperationBase &self, PyOperationBase &other) {
              return &self.getOperation() == &other.getOperation();
            })
       .def("__eq__",
            [](PyOperationBase &self, py::object other) { return false; })
       .def_property_readonly("attributes",
                              [](PyOperationBase &self) {
py::class_<PyOperationBase>(m, "_OperationBase", py::module_local())
       .def(
           "verify",
           [](PyOperationBase &self) {
             return mlirOperationVerify(self.getOperation());
           },
           "Verify the operation and return true if it passes, false if it "
           "fails.");
 
-  py::class_<PyOperation, PyOperationBase>(m, "Operation")
+  py::class_<PyOperation, PyOperationBase>(m, "Operation", py::module_local())
       .def_static("create", &PyOperation::create, py::arg("name"),
                   py::arg("results") = py::none(),
                   py::arg("operands") = py::none(),
                   py::arg("attributes") = py::none(),
                   py::arg("successors") = py::none(), py::arg("regions") = 0,
                   py::arg("loc") = py::none(), py::arg("ip") = py::none(),
                   kOperationCreateDocstring)
       .def_property_readonly("parent",
py::class_<PyOperation, PyOperationBase>(m, "Operation", py::module_local())
           [](PyOperation &self) {
             self.checkValid();
             return self.getContext().getObject();
           },
           "Context that owns the Operation")
       .def_property_readonly("opview", &PyOperation::createOpView);
 
   auto opViewClass =
-      py::class_<PyOpView, PyOperationBase>(m, "OpView")
+      py::class_<PyOpView, PyOperationBase>(m, "OpView", py::module_local())
           .def(py::init<py::object>())
           .def_property_readonly("operation", &PyOpView::getOperationObject)
           .def_property_readonly(
               "context",
               [](PyOpView &self) {
                 return self.getOperation().getContext().getObject();
               },
               "Context that owns the Operation")
           .def("__str__", [](PyOpView &self) {
             return py::str(self.getOperationObject());
           });
   opViewClass.attr("_ODS_REGIONS") = py::make_tuple(0, true);
   opViewClass.attr("_ODS_OPERAND_SEGMENTS") = py::none();
   opViewClass.attr("_ODS_RESULT_SEGMENTS") = py::none();
   opViewClass.attr("build_generic") = classmethod(
       &PyOpView::buildGeneric, py::arg("cls"), py::arg("results") = py::none(),
       py::arg("operands") = py::none(), py::arg("attributes") = py::none(),
       py::arg("successors") = py::none(), py::arg("regions") = py::none(),
       py::arg("loc") = py::none(), py::arg("ip") = py::none(),
       "Builds a specific, generated OpView based on class level attributes.");
 
   //----------------------------------------------------------------------------
   // Mapping of PyRegion.
   //----------------------------------------------------------------------------
-  py::class_<PyRegion>(m, "Region")
+  py::class_<PyRegion>(m, "Region", py::module_local())
       .def_property_readonly(
           "blocks",
           [](PyRegion &self) {
             return PyBlockList(self.getParentOperation(), self.get());
           },
           "Returns a forward-optimized sequence of blocks.")
       .def(
           "__iter__",
           [](PyRegion &self) {
             self.checkValid();
             MlirBlock firstBlock = mlirRegionGetFirstBlock(self.get());
             return PyBlockIterator(self.getParentOperation(), firstBlock);
           },
           "Iterates over blocks in the region.")
       .def("__eq__",
            [](PyRegion &self, PyRegion &other) {
              return self.get().ptr == other.get().ptr;
            })
       .def("__eq__", [](PyRegion &self, py::object &other) { return false; });
 
   //----------------------------------------------------------------------------
   // Mapping of PyBlock.
   //----------------------------------------------------------------------------
-  py::class_<PyBlock>(m, "Block")
+  py::class_<PyBlock>(m, "Block", py::module_local())
       .def_property_readonly(
           "owner",
           [](PyBlock &self) {
             return self.getParentOperation()->createOpView();
           },
           "Returns the owning operation of this block.")
       .def_property_readonly(
           "arguments",
py::class_<PyBlock>(m, "Block", py::module_local())
             return printAccum.join();
           },
           "Returns the assembly form of the block.");
 
   //----------------------------------------------------------------------------
   // Mapping of PyInsertionPoint.
   //----------------------------------------------------------------------------
 
-  py::class_<PyInsertionPoint>(m, "InsertionPoint")
+  py::class_<PyInsertionPoint>(m, "InsertionPoint", py::module_local())
       .def(py::init<PyBlock &>(), py::arg("block"),
            "Inserts after the last operation but still inside the block.")
       .def("__enter__", &PyInsertionPoint::contextEnter)
       .def("__exit__", &PyInsertionPoint::contextExit)
       .def_property_readonly_static(
           "current",
           [](py::object & /*class*/) {
             auto *ip = PyThreadContextEntry::getDefaultInsertionPoint();
py::class_<PyInsertionPoint>(m, "InsertionPoint", py::module_local())
       .def_static("at_block_terminator", &PyInsertionPoint::atBlockTerminator,
                   py::arg("block"), "Inserts before the block terminator.")
       .def("insert", &PyInsertionPoint::insert, py::arg("operation"),
            "Inserts an operation.");
 
   //----------------------------------------------------------------------------
   // Mapping of PyAttribute.
   //----------------------------------------------------------------------------
-  py::class_<PyAttribute>(m, "Attribute")
+  py::class_<PyAttribute>(m, "Attribute", py::module_local())
       // Delegate to the PyAttribute copy constructor, which will also lifetime
       // extend the backing context which owns the MlirAttribute.
       .def(py::init<PyAttribute &>(), py::arg("cast_from_type"),
            "Casts the passed attribute to the generic Attribute")
       .def_property_readonly(MLIR_PYTHON_CAPI_PTR_ATTR,
                              &PyAttribute::getCapsule)
       .def(MLIR_PYTHON_CAPI_FACTORY_ATTR, &PyAttribute::createFromCapsule)
       .def_static(
py::class_<PyAttribute>(m, "Attribute", py::module_local())
                            printAccum.getUserData());
         printAccum.parts.append(")");
         return printAccum.join();
       });
 
   //----------------------------------------------------------------------------
   // Mapping of PyNamedAttribute
   //----------------------------------------------------------------------------
-  py::class_<PyNamedAttribute>(m, "NamedAttribute")
+  py::class_<PyNamedAttribute>(m, "NamedAttribute", py::module_local())
       .def("__repr__",
            [](PyNamedAttribute &self) {
              PyPrintAccumulator printAccum;
              printAccum.parts.append("NamedAttribute(");
              printAccum.parts.append(
                  mlirIdentifierStr(self.namedAttr.name).data);
              printAccum.parts.append("=");
              mlirAttributePrint(self.namedAttr.attribute,
py::class_<PyNamedAttribute>(m, "NamedAttribute", py::module_local())
             return PyAttribute(std::move(contextRef), self.namedAttr.attribute);
           },
           py::keep_alive<0, 1>(),
           "The underlying generic attribute of the NamedAttribute binding");
 
   //----------------------------------------------------------------------------
   // Mapping of PyType.
   //----------------------------------------------------------------------------
-  py::class_<PyType>(m, "Type")
+  py::class_<PyType>(m, "Type", py::module_local())
       // Delegate to the PyType copy constructor, which will also lifetime
       // extend the backing context which owns the MlirType.
       .def(py::init<PyType &>(), py::arg("cast_from_type"),
            "Casts the passed type to the generic Type")
       .def_property_readonly(MLIR_PYTHON_CAPI_PTR_ATTR, &PyType::getCapsule)
       .def(MLIR_PYTHON_CAPI_FACTORY_ATTR, &PyType::createFromCapsule)
       .def_static(
           "parse",
py::class_<PyType>(m, "Type", py::module_local())
         mlirTypePrint(self, printAccum.getCallback(), printAccum.getUserData());
         printAccum.parts.append(")");
         return printAccum.join();
       });
 
   //----------------------------------------------------------------------------
   // Mapping of Value.
   //----------------------------------------------------------------------------
-  py::class_<PyValue>(m, "Value")
+  py::class_<PyValue>(m, "Value", py::module_local())
       .def_property_readonly(MLIR_PYTHON_CAPI_PTR_ATTR, &PyValue::getCapsule)
       .def(MLIR_PYTHON_CAPI_FACTORY_ATTR, &PyValue::createFromCapsule)
       .def_property_readonly(
           "context",
           [](PyValue &self) { return self.getParentOperation()->getContext(); },
           "Context in which the value lives.")
       .def(
           "dump", [](PyValue &self) { mlirValueDump(self.get()); },

mlir/lib/Bindings/Python/MainModule.cpp

 
 // -----------------------------------------------------------------------------
 // Module initialization.
 // -----------------------------------------------------------------------------
 
 PYBIND11_MODULE(_mlir, m) {
   m.doc() = "MLIR Python Native Extension";
 
-  py::class_<PyGlobals>(m, "_Globals")
+  py::class_<PyGlobals>(m, "_Globals", py::module_local())
       .def_property("dialect_search_modules",
                     &PyGlobals::getDialectSearchPrefixes,
                     &PyGlobals::setDialectSearchPrefixes)
       .def("append_dialect_search_prefix",
            [](PyGlobals &self, std::string moduleName) {
              self.getDialectSearchPrefixes().push_back(std::move(moduleName));
              self.clearImportCache();
            })

mlir/lib/Bindings/Python/Pass.cpp

 
 } // anonymous namespace
 
 /// Create the `mlir.passmanager` here.
 void mlir::python::populatePassManagerSubmodule(py::module &m) {
   //----------------------------------------------------------------------------
   // Mapping of the top-level PassManager
   //----------------------------------------------------------------------------
-  py::class_<PyPassManager>(m, "PassManager")
+  py::class_<PyPassManager>(m, "PassManager", py::module_local())
       .def(py::init<>([](DefaultingPyMlirContext context) {
              MlirPassManager passManager =
                  mlirPassManagerCreate(context->get());
              return new PyPassManager(passManager);
            }),
            py::arg("context") = py::none(),
            "Create a new PassManager for the current (or provided) Context.")
       .def_property_readonly(MLIR_PYTHON_CAPI_PTR_ATTR,