[mlir][Python] Add missing capsule->module and Context.create_module.

[stellaraccident]

Summary

• Discovered missing in npcomp.

[mehdi_amini]

mlir/include/mlir-c/Bindings/Python/Interop.h

93

This function isn't used anywhere right now?

[stellaraccident]

mlir/include/mlir-c/Bindings/Python/Interop.h

93

Good catch. I should implement the module uniquing like I did with context, then it will be used and testable easily. We probably want that anyway, so it would be better than creating a hidden API just for testing this.

[stellaraccident]

Intern modules and add tests.

[stellaraccident]

PTAL - with module interning, the capsule accessors are verified internally via tests now.

[mehdi_amini]

mlir/lib/Bindings/Python/IRModules.cpp

594

Don't you need to take the GIL?

(seeing that you took it line 602)

[stellaraccident]

Rebase and comment.

[mehdi_amini]

mlir/lib/Bindings/Python/IRModules.cpp

594

You haven't answered here, so I am not sure if you missed it or not?

mlir/include/mlir-c/Bindings/Python/Interop.h

 #ifdef __cplusplus
   void *ptr = const_cast<void *>(module.ptr);
 #else
   void *ptr = (void *)ptr;
 #endif
   return PyCapsule_New(ptr, MLIR_PYTHON_CAPSULE_MODULE, NULL);
 }
 
+/** Extracts an MlirModule from a capsule as produced from
+ * mlirPythonModuleToCapsule. If the capsule is not of the right type, then
+ * a null module is returned (as checked via mlirModuleIsNull). In such a
+ * case, the Python APIs will have already set an error. */
+inline MlirModule mlirPythonCapsuleToModule(PyObject *capsule) {
+  void *ptr = PyCapsule_GetPointer(capsule, MLIR_PYTHON_CAPSULE_MODULE);
+  MlirModule module = {ptr};
+  return module;
+}
+
 #ifdef __cplusplus
 }
 #endif
 
 #endif // MLIR_C_BINDINGS_PYTHON_INTEROP_H

mlir/lib/Bindings/Python/IRModules.h

PyMlirContextRef getRef() {
     return PyMlirContextRef(this, pybind11::cast(this));
   }
 
   /// Gets a capsule wrapping the void* within the MlirContext.
   pybind11::object getCapsule();
 
   /// Creates a PyMlirContext from the MlirContext wrapped by a capsule.
   /// Note that PyMlirContext instances are uniqued, so the returned object
-  /// may be a pre-existing object.
+  /// may be a pre-existing object. Ownership of the underlying MlirContext
+  /// is taken by calling this function.
   static pybind11::object createFromCapsule(pybind11::object capsule);
 
   /// Gets the count of live context objects. Used for testing.
   static size_t getLiveCount();
 
   /// Gets the count of live operations associated with this context.
   /// Used for testing.
   size_t getLiveOperationCount();
 
+  /// Gets the count of live modules associated with this context.
+  /// Used for testing.
+  size_t getLiveModuleCount();
+
   /// Creates an operation. See corresponding python docstring.
   pybind11::object
   createOperation(std::string name, PyLocation location,
                   llvm::Optional<std::vector<PyType *>> results,
                   llvm::Optional<pybind11::dict> attributes,
                   llvm::Optional<std::vector<PyBlock *>> successors,
                   int regions);
 
 private:
   PyMlirContext(MlirContext context);
   // Interns the mapping of live MlirContext::ptr to PyMlirContext instances,
   // preserving the relationship that an MlirContext maps to a single
   // PyMlirContext wrapper. This could be replaced in the future with an
   // extension mechanism on the MlirContext for stashing user pointers.
   // Note that this holds a handle, which does not imply ownership.
   // Mappings will be removed when the context is destructed.
   using LiveContextMap = llvm::DenseMap<void *, PyMlirContext *>;
   static LiveContextMap &getLiveContexts();
 
+  // Interns all live modules associated with this context. Modules tracked
+  // in this map are valid. When a module is invalidated, it is removed
+  // from this map, and while it still exists as an instance, any
+  // attempt to access it will raise an error.
+  using LiveModuleMap =
+      llvm::DenseMap<const void *, std::pair<pybind11::handle, PyModule *>>;
+  LiveModuleMap liveModules;
+
   // Interns all live operations associated with this context. Operations
   // tracked in this map are valid. When an operation is invalidated, it is
   // removed from this map, and while it still exists as an instance, any
   // attempt to access it will raise an error.
   using LiveOperationMap =
       llvm::DenseMap<void *, std::pair<pybind11::handle, PyOperation *>>;
   LiveOperationMap liveOperations;
 
   MlirContext context;
+  friend class PyModule;
   friend class PyOperation;
 };
 
 /// Base class for all objects that directly or indirectly depend on an
 /// MlirContext. The lifetime of the context will extend at least to the
 /// lifetime of these instances.
 /// Immutable objects that depend on a context extend this directly.
 class BaseContextObject {

 };
 
 /// Wrapper around MlirModule.
 /// This is the top-level, user-owned object that contains regions/ops/blocks.
 class PyModule;
 using PyModuleRef = PyObjectRef<PyModule>;
 class PyModule : public BaseContextObject {
 public:
-  /// Creates a reference to the module
-  static PyModuleRef create(PyMlirContextRef contextRef, MlirModule module);
+  /// Returns a PyModule reference for the given MlirModule. This may return
+  /// a pre-existing or new object.
+  static PyModuleRef forModule(MlirModule module);
   PyModule(PyModule &) = delete;
-  ~PyModule() {
-    if (module.ptr)
-      mlirModuleDestroy(module);
-  }
+  PyModule(PyMlirContext &&) = delete;
+  ~PyModule();
 
   /// Gets the backing MlirModule.
   MlirModule get() { return module; }
 
   /// Gets a strong reference to this module.
   PyModuleRef getRef() {
     return PyModuleRef(this,
                        pybind11::reinterpret_borrow<pybind11::object>(handle));
   }
 
   /// Gets a capsule wrapping the void* within the MlirModule.
   /// Note that the module does not (yet) provide a corresponding factory for
   /// constructing from a capsule as that would require uniquing PyModule
   /// instances, which is not currently done.
   pybind11::object getCapsule();
 
+  /// Creates a PyModule from the MlirModule wrapped by a capsule.
+  /// Note that PyModule instances are uniqued, so the returned object
+  /// may be a pre-existing object. Ownership of the underlying MlirModule
+  /// is taken by calling this function.
+  static pybind11::object createFromCapsule(pybind11::object capsule);
+
 private:
-  PyModule(PyMlirContextRef contextRef, MlirModule module)
-      : BaseContextObject(std::move(contextRef)), module(module) {}
+  PyModule(PyMlirContextRef contextRef, MlirModule module);
   MlirModule module;
   pybind11::handle handle;
 };
 
 /// Wrapper around PyOperation.
 /// Operations exist in either an attached (dependent) or detached (top-level)
 /// state. In the detached state (as on creation), an operation is owned by
 /// the creator and its lifetime extends either until its reference count

mlir/lib/Bindings/Python/IRModules.cpp

PyMlirContext::LiveContextMap &PyMlirContext::getLiveContexts() {
   static LiveContextMap liveContexts;
   return liveContexts;
 }
 
 size_t PyMlirContext::getLiveCount() { return getLiveContexts().size(); }
 
 size_t PyMlirContext::getLiveOperationCount() { return liveOperations.size(); }
 
+size_t PyMlirContext::getLiveModuleCount() { return liveModules.size(); }
+
 py::object PyMlirContext::createOperation(
     std::string name, PyLocation location,
     llvm::Optional<std::vector<PyType *>> results,
     llvm::Optional<py::dict> attributes,
     llvm::Optional<std::vector<PyBlock *>> successors, int regions) {
   llvm::SmallVector<MlirType, 4> mlirResults;
   llvm::SmallVector<MlirBlock, 4> mlirSuccessors;
   llvm::SmallVector<std::pair<std::string, MlirAttribute>, 4> mlirAttributes;
py::object PyMlirContext::createOperation(
   MlirOperation operation = mlirOperationCreate(&state);
   return PyOperation::createDetached(getRef(), operation).releaseObject();
 }
 
 //------------------------------------------------------------------------------
 // PyModule
 //------------------------------------------------------------------------------
 
-PyModuleRef PyModule::create(PyMlirContextRef contextRef, MlirModule module) {
+PyModule::PyModule(PyMlirContextRef contextRef, MlirModule module)
+    : BaseContextObject(std::move(contextRef)), module(module) {}
+
+PyModule::~PyModule() {
+  py::gil_scoped_acquire acquire;
+  auto &liveModules = getContext()->liveModules;
+  assert(liveModules.count(module.ptr) == 1 &&
+         "destroying module not in live map");
+  liveModules.erase(module.ptr);
+  mlirModuleDestroy(module);
+}
+
+PyModuleRef PyModule::forModule(MlirModule module) {
+  MlirContext context = mlirModuleGetContext(module);
+  PyMlirContextRef contextRef = PyMlirContext::forContext(context);
+
+  py::gil_scoped_acquire acquire;
+  auto &liveModules = contextRef->liveModules;
+  auto it = liveModules.find(module.ptr);
+  if (it == liveModules.end()) {
+    // Create.
     PyModule *unownedModule = new PyModule(std::move(contextRef), module);
     // Note that the default return value policy on cast is automatic_reference,
     // which does not take ownership (delete will not be called).
     // Just be explicit.
     py::object pyRef =
         py::cast(unownedModule, py::return_value_policy::take_ownership);
     unownedModule->handle = pyRef;
+    liveModules[module.ptr] =
+        std::make_pair(unownedModule->handle, unownedModule);
     return PyModuleRef(unownedModule, std::move(pyRef));
   }
+  // Use existing.
+  PyModule *existing = it->second.second;
+  py::object pyRef = py::reinterpret_borrow<py::object>(it->second.first);
+  return PyModuleRef(existing, std::move(pyRef));
+}
+
+py::object PyModule::createFromCapsule(py::object capsule) {
+  MlirModule rawModule = mlirPythonCapsuleToModule(capsule.ptr());
+  if (mlirModuleIsNull(rawModule))
+    throw py::error_already_set();
+  return forModule(rawModule).releaseObject();
+}
 
 py::object PyModule::getCapsule() {
   return py::reinterpret_steal<py::object>(mlirPythonModuleToCapsule(get()));
 }
 
 //------------------------------------------------------------------------------
 // PyOperation
 //------------------------------------------------------------------------------
py::class_<PyMlirContext>(m, "Context")
       .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)
+      .def("_get_live_module_count", &PyMlirContext::getLiveModuleCount)
       .def_property_readonly(MLIR_PYTHON_CAPI_PTR_ATTR,
                              &PyMlirContext::getCapsule)
       .def(MLIR_PYTHON_CAPI_FACTORY_ATTR, &PyMlirContext::createFromCapsule)
       .def_property(
           "allow_unregistered_dialects",
           [](PyMlirContext &self) -> bool {
             return mlirContextGetAllowUnregisteredDialects(self.get());
           },
py::class_<PyMlirContext>(m, "Context")
                 mlirModuleCreateParse(self.get(), moduleAsm.c_str());
             // TODO: Rework error reporting once diagnostic engine is exposed
             // in C API.
             if (mlirModuleIsNull(module)) {
               throw SetPyError(
                   PyExc_ValueError,
                   "Unable to parse module assembly (see diagnostics)");
             }
-            return PyModule::create(self.getRef(), module).releaseObject();
+            return PyModule::forModule(module).releaseObject();
           },
           kContextParseDocstring)
       .def(
+          "create_module",
+          [](PyMlirContext &self, PyLocation &loc) {
+            MlirModule module = mlirModuleCreateEmpty(loc.loc);
+            return PyModule::forModule(module).releaseObject();
+          },
+          py::arg("loc"), "Creates an empty module")
+      .def(
           "parse_attr",
           [](PyMlirContext &self, std::string attrSpec) {
             MlirAttribute type =
                 mlirAttributeParseGet(self.get(), attrSpec.c_str());
             // TODO: Rework error reporting once diagnostic engine is exposed
             // in C API.
             if (mlirAttributeIsNull(type)) {
               throw SetPyError(PyExc_ValueError,
py::class_<PyMlirContext>(m, "Context")
           [](PyMlirContext &self, std::string filename, int line, int col) {
             return PyLocation(self.getRef(),
                               mlirLocationFileLineColGet(
                                   self.get(), filename.c_str(), line, col));
           },
           kContextGetFileLocationDocstring, py::arg("filename"),
           py::arg("line"), py::arg("col"));
 
-  py::class_<PyLocation>(m, "Location").def("__repr__", [](PyLocation &self) {
+  py::class_<PyLocation>(m, "Location")
+      .def_property_readonly(
+          "context",
+          [](PyLocation &self) { return self.getContext().getObject(); },
+          "Context that owns the Location")
+      .def("__repr__", [](PyLocation &self) {
         PyPrintAccumulator printAccum;
         mlirLocationPrint(self.loc, printAccum.getCallback(),
                           printAccum.getUserData());
         return printAccum.join();
       });
 
   // Mapping of Module
   py::class_<PyModule>(m, "Module")
       .def_property_readonly(MLIR_PYTHON_CAPI_PTR_ATTR, &PyModule::getCapsule)
+      .def(MLIR_PYTHON_CAPI_FACTORY_ATTR, &PyModule::createFromCapsule)
+      .def_property_readonly(
+          "context",
+          [](PyModule &self) { return self.getContext().getObject(); },
+          "Context that created the Module")
       .def_property_readonly(
           "operation",
           [](PyModule &self) {
             return PyOperation::forOperation(self.getContext(),
                                              mlirModuleGetOperation(self.get()),
                                              self.getRef().releaseObject())
                 .releaseObject();
           },
py::class_<PyModule>(m, "Module")
                                printAccum.getUserData());
             return printAccum.join();
           },
           kOperationStrDunderDocstring);
 
   // Mapping of Operation.
   py::class_<PyOperation>(m, "Operation")
       .def_property_readonly(
+          "context",
+          [](PyOperation &self) { return self.getContext().getObject(); },
+          "Context that owns the Operation")
+      .def_property_readonly(
           "regions",
           [](PyOperation &self) { return PyRegionList(self.getRef()); })
       .def("__iter__",
            [](PyOperation &self) { return PyRegionIterator(self.getRef()); })
       .def(
           "__str__",
           [](PyOperation &self) {
             self.checkValid();
py::class_<PyBlock>(m, "Block")
             mlirBlockPrint(self.get(), printAccum.getCallback(),
                            printAccum.getUserData());
             return printAccum.join();
           },
           kTypeStrDunderDocstring);
 
   // Mapping of Type.
   py::class_<PyAttribute>(m, "Attribute")
+      .def_property_readonly(
+          "context",
+          [](PyAttribute &self) { return self.getContext().getObject(); },
+          "Context that owns the Attribute")
       .def(
           "get_named",
           [](PyAttribute &self, std::string name) {
             return PyNamedAttribute(self.attr, std::move(name));
           },
           py::keep_alive<0, 1>(), "Binds a name to the attribute")
       .def("__eq__",
            [](PyAttribute &self, py::object &other) {
void mlir::python::populateIRSubmodule(py::module &m) {
   // Standard attribute bindings.
   PyFloatAttribute::bind(m);
   PyIntegerAttribute::bind(m);
   PyBoolAttribute::bind(m);
   PyStringAttribute::bind(m);
 
   // Mapping of Type.
   py::class_<PyType>(m, "Type")
+      .def_property_readonly(
+          "context", [](PyType &self) { return self.getContext().getObject(); },
+          "Context that owns the Type")
       .def("__eq__",
            [](PyType &self, py::object &other) {
              try {
                PyType otherType = other.cast<PyType>();
                return self == otherType;
              } catch (std::exception &e) {
                return false;
              }

mlir/test/Bindings/Python/ir_attributes.py

 # RUN: %PYTHON %s | FileCheck %s
 
 import gc
 import mlir
 
 def run(f):
   print("\nTEST:", f.__name__)
   f()
   gc.collect()
   assert mlir.ir.Context._get_live_count() == 0
 
 
 # CHECK-LABEL: TEST: testParsePrint
 def testParsePrint():
   ctx = mlir.ir.Context()
   t = ctx.parse_attr('"hello"')
+  assert t.context is ctx
   ctx = None
   gc.collect()
   # CHECK: "hello"
   print(str(t))
   # CHECK: Attribute("hello")
   print(repr(t))
 
 run(testParsePrint)

mlir/test/Bindings/Python/ir_location.py

 # RUN: %PYTHON %s | FileCheck %s
 
 import gc
 import mlir
 
 def run(f):
   print("\nTEST:", f.__name__)
   f()
   gc.collect()
   assert mlir.ir.Context._get_live_count() == 0
 
 
 # CHECK-LABEL: TEST: testUnknown
 def testUnknown():
   ctx = mlir.ir.Context()
   loc = ctx.get_unknown_location()
+  assert loc.context is ctx
   ctx = None
   gc.collect()
   # CHECK: unknown str: loc(unknown)
   print("unknown str:", str(loc))
   # CHECK: unknown repr: loc(unknown)
   print("unknown repr:", repr(loc))
 
 run(testUnknown)

mlir/test/Bindings/Python/ir_module.py

 
 
 # Verify successful parse.
 # CHECK-LABEL: TEST: testParseSuccess
 # CHECK: module @successfulParse
 def testParseSuccess():
   ctx = mlir.ir.Context()
   module = ctx.parse_module(r"""module @successfulParse {}""")
+  assert module.context is ctx
   print("CLEAR CONTEXT")
   ctx = None  # Ensure that module captures the context.
   gc.collect()
   module.dump()  # Just outputs to stderr. Verifies that it functions.
   print(str(module))
 
 run(testParseSuccess)
 
 
 # Verify parse error.
 # CHECK-LABEL: TEST: testParseError
 # CHECK: testParseError: Unable to parse module assembly (see diagnostics)
 def testParseError():
   ctx = mlir.ir.Context()
   try:
     module = ctx.parse_module(r"""}SYNTAX ERROR{""")
   except ValueError as e:
     print("testParseError:", e)
   else:
     print("Exception not produced")
 
 run(testParseError)
 
 
+# Verify successful parse.
+# CHECK-LABEL: TEST: testCreateEmpty
+# CHECK: module {
+def testCreateEmpty():
+  ctx = mlir.ir.Context()
+  loc = ctx.get_unknown_location()
+  module = ctx.create_module(loc)
+  print("CLEAR CONTEXT")
+  ctx = None  # Ensure that module captures the context.
+  gc.collect()
+  print(str(module))
+
+run(testCreateEmpty)
+
+
 # Verify round-trip of ASM that contains unicode.
 # Note that this does not test that the print path converts unicode properly
 # because MLIR asm always normalizes it to the hex encoding.
 # CHECK-LABEL: TEST: testRoundtripUnicode
 # CHECK: func @roundtripUnicode()
 # CHECK: foo = "\F0\9F\98\8A"
 def testRoundtripUnicode():
   ctx = mlir.ir.Context()
   module = ctx.parse_module(r"""
     func @roundtripUnicode() attributes { foo = "😊" }
   """)
   print(str(module))
 
 run(testRoundtripUnicode)
 
 
 # Tests that module.operation works and correctly interns instances.
 # CHECK-LABEL: TEST: testModuleOperation
 def testModuleOperation():
   ctx = mlir.ir.Context()
   module = ctx.parse_module(r"""module @successfulParse {}""")
+  assert ctx._get_live_module_count() == 1
   op1 = module.operation
   assert ctx._get_live_operation_count() == 1
   # CHECK: module @successfulParse
   print(op1)
 
   # Ensure that operations are the same on multiple calls.
   op2 = module.operation
   assert ctx._get_live_operation_count() == 1
   assert op1 is op2
 
   # Ensure that if module is de-referenced, the operations are still valid.
   module = None
   gc.collect()
   print(op1)
 
   # Collect and verify lifetime.
   op1 = None
   op2 = None
   gc.collect()
   print("LIVE OPERATIONS:", ctx._get_live_operation_count())
   assert ctx._get_live_operation_count() == 0
+  assert ctx._get_live_module_count() == 0
 
 run(testModuleOperation)
 
 
 # CHECK-LABEL: TEST: testModuleCapsule
 def testModuleCapsule():
   ctx = mlir.ir.Context()
   module = ctx.parse_module(r"""module @successfulParse {}""")
+  assert ctx._get_live_module_count() == 1
   # CHECK: "mlir.ir.Module._CAPIPtr"
-  print(module._CAPIPtr)
+  module_capsule = module._CAPIPtr
+  print(module_capsule)
+  module_dup = mlir.ir.Module._CAPICreate(module_capsule)
+  assert module is module_dup
+  assert module_dup.context is ctx
+  # Gc and verify destructed.
+  module = None
+  module_capsule = None
+  module_dup = None
+  gc.collect()
+  assert ctx._get_live_module_count() == 0
+
 
 run(testModuleCapsule)

mlir/test/Bindings/Python/ir_operation.py

def testTraverseOpRegionBlockIterators():
   ctx.allow_unregistered_dialects = True
   module = ctx.parse_module(r"""
     func @f1(%arg0: i32) -> i32 {
       %1 = "custom.addi"(%arg0, %arg0) : (i32, i32) -> i32
       return %1 : i32
     }
   """)
   op = module.operation
+  assert op.context is ctx
   # Get the block using iterators off of the named collections.
   regions = list(op.regions)
   blocks = list(regions[0].blocks)
   # CHECK: MODULE REGIONS=1 BLOCKS=1
   print(f"MODULE REGIONS={len(regions)} BLOCKS={len(blocks)}")
 
   # Get the regions and blocks from the default collections.
   default_regions = list(op)

mlir/test/Bindings/Python/ir_types.py

 # RUN: %PYTHON %s | FileCheck %s
 
 import gc
 import mlir
 
 def run(f):
   print("\nTEST:", f.__name__)
   f()
   gc.collect()
   assert mlir.ir.Context._get_live_count() == 0
 
 
 # CHECK-LABEL: TEST: testParsePrint
 def testParsePrint():
   ctx = mlir.ir.Context()
   t = ctx.parse_type("i32")
+  assert t.context is ctx
   ctx = None
   gc.collect()
   # CHECK: i32
   print(str(t))
   # CHECK: Type(i32)
   print(repr(t))
 
 run(testParsePrint)