[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

600

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

600

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

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

 |*                                                                            *|
 |* This header is standalone, requiring nothing beyond normal linking against *|
 |* the Python implementation.                                                 *|
 \*===----------------------------------------------------------------------===*/
 
 #ifndef MLIR_C_BINDINGS_PYTHON_INTEROP_H
 #define MLIR_C_BINDINGS_PYTHON_INTEROP_H
 
 #include <Python.h>
 
 #include "mlir-c/IR.h"
 
 #define MLIR_PYTHON_CAPSULE_CONTEXT "mlir.ir.Context._CAPIPtr"
 #define MLIR_PYTHON_CAPSULE_MODULE "mlir.ir.Module._CAPIPtr"
 
 /** Attribute on MLIR Python objects that expose their C-API pointer.
  * This will be a type-specific capsule created as per one of the helpers

 #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.cpp

PyModuleRef PyModule::create(PyMlirContextRef contextRef, MlirModule module) {
   unownedModule->handle = pyRef;
   return PyModuleRef(unownedModule, std::move(pyRef));
 }
 
 py::object PyModule::getCapsule() {
   return py::reinterpret_steal<py::object>(mlirPythonModuleToCapsule(get()));
 }
 
 //------------------------------------------------------------------------------
 // PyOperation
 //------------------------------------------------------------------------------
 
 PyOperation::PyOperation(PyMlirContextRef contextRef, MlirOperation operation)
     : BaseContextObject(std::move(contextRef)), operation(operation) {}
 
 PyOperation::~PyOperation() {
   auto &liveOperations = getContext()->liveOperations;
py::class_<PyMlirContext>(m, "Context")
               throw SetPyError(
                   PyExc_ValueError,
                   "Unable to parse module assembly (see diagnostics)");
             }
             return PyModule::create(self.getRef(), module).releaseObject();
           },
           kContextParseDocstring)
       .def(
+          "create_module",
+          [](PyMlirContext &self, PyLocation &loc) {
+            MlirModule module = mlirModuleCreateEmpty(loc.loc);
+            return PyModule::create(self.getRef(), 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,

mlir/test/Bindings/Python/ir_module.py

def testParseError():
   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()