[mlir][Python] Custom python op view wrappers for building and traversing.

[stellaraccident]

Summary

• Still rough edges that need more sugar but the bones are there. Notes left in the test case for things that can be improved.
• Does not actually yield custom OpViews yet for traversing. Will rework that in a followup.

[stellaraccident]

Respond to conversation.

[stellaraccident]

Send for review

[ftynse]

mlir/lib/Bindings/Python/Globals.h

58 ↗

(On Diff #300875)

Nit: if the mapping already exists

64 ↗

(On Diff #300875)

Same nit

71 ↗

(On Diff #300875)

Nit: "an the"

81 ↗

(On Diff #300875)

What is the convention for bool values? true = loaded, false = not found, no-value = not attempted?

mlir/lib/Bindings/Python/IRModules.h

454

Is this change intentional?

mlir/lib/Bindings/Python/MainModule.cpp

57 ↗

(On Diff #300875)

Can this ever contain false? If not, we rather want a StringSet<> for this.

[stellaraccident]

mlir/lib/Bindings/Python/Globals.h

81 ↗

(On Diff #300875)

Changed to StringSet.

mlir/lib/Bindings/Python/IRModules.h

454

Thanks from a previous iteration of the patch where I had to capture the module in the method, but since I've pulled this out to the PyGlobals now, it is unnecessary. Reverted.

mlir/lib/Bindings/Python/MainModule.cpp

57 ↗

(On Diff #300875)

Thanks. Originally was mapping to the module, removed the mapped value and should have dropped it back to a set then. Done.

[stellaraccident]

Address comments.

mlir/lib/Bindings/Python/CMakeLists.txt

 ################################################################################
 # Copy python source tree.
 ################################################################################
 
 set(PY_SRC_FILES
   mlir/__init__.py
+  mlir/extensions/__init__.py
+  mlir/extensions/std_dialect.py
 )
 
 add_custom_target(MLIRBindingsPythonSources ALL
   DEPENDS ${PY_SRC_FILES}
 )
 
 foreach(PY_SRC_FILE ${PY_SRC_FILES})
   set(PY_DEST_FILE "${PROJECT_BINARY_DIR}/python/${PY_SRC_FILE}")

mlir/lib/Bindings/Python/IRModules.h

 //===- IRModules.h - IR Submodules of pybind module -----------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_BINDINGS_PYTHON_IRMODULES_H
 #define MLIR_BINDINGS_PYTHON_IRMODULES_H
 
 #include <pybind11/pybind11.h>
 
 #include "mlir-c/IR.h"
 #include "llvm/ADT/DenseMap.h"
 
 namespace mlir {
 namespace python {
 
 class PyBlock;
public:
 
   /// Accesses the context reference.
   PyMlirContextRef &getContext() { return contextRef; }
 
 private:
   PyMlirContextRef contextRef;
 };
 
+/// Wrapper around an MlirDialect. This is exported as `DialectDescriptor` in
+/// order to differentiate it from the `Dialect` base class which is extended by
+/// plugins which extend dialect functionality through extension python code.
+/// This should be seen as the "low-level" object and `Dialect` as the
+/// high-level, user facing object.
+class PyDialectDescriptor : public BaseContextObject {
+public:
+  PyDialectDescriptor(PyMlirContextRef contextRef, MlirDialect dialect)
+      : BaseContextObject(std::move(contextRef)), dialect(dialect) {}
+
+  MlirDialect get() { return dialect; }
+
+private:
+  MlirDialect dialect;
+};
+
+/// User-level object for accessing dialects with dotted syntax such as:
+///   ctx.dialect.std
+class PyDialects : public BaseContextObject {
+public:
+  PyDialects(PyMlirContextRef contextRef)
+      : BaseContextObject(std::move(contextRef)) {}
+
+  MlirDialect getDialectForKey(const std::string &key, bool attrError);
+};
+
+/// User-level dialect object. For dialects that have a registered extension,
+/// this will be the base class of the extension dialect type. For un-extended,
+/// objects of this type will be returned directly.
+class PyDialect {
+public:
+  PyDialect(pybind11::object descriptor) : descriptor(std::move(descriptor)) {}
+
+  pybind11::object getDescriptor() { return descriptor; }
+
+private:
+  pybind11::object descriptor;
+};
+
 /// Wrapper around an MlirLocation.
 class PyLocation : public BaseContextObject {
 public:
   PyLocation(PyMlirContextRef contextRef, MlirLocation loc)
       : BaseContextObject(std::move(contextRef)), loc(loc) {}
   MlirLocation loc;
 };
 
public:
 
   void checkValid() { return parentOperation->checkValid(); }
 
 private:
   PyOperationRef parentOperation;
   MlirValue value;
 };
 
-void populateIRSubmodule(pybind11::module &m);
+void populateIRSubmodule(pybind11::module m);
 
 } // namespace python
 } // namespace mlir
 
 #endif // MLIR_BINDINGS_PYTHON_IRMODULES_H

mlir/lib/Bindings/Python/IRModules.cpp

mlirOperationStateAddOwnedRegions(&state, mlirRegions.size(),
                                       mlirRegions.data());
   }
 
   // Construct the operation.
   MlirOperation operation = mlirOperationCreate(&state);
   return PyOperation::createDetached(getRef(), operation).releaseObject();
 }
 
+MlirDialect PyDialects::getDialectForKey(const std::string &key,
+                                         bool attrError) {
+  MlirDialect dialect = mlirContextGetOrLoadDialect(getContext()->get(),
+                                                    {key.data(), key.size()});
+  if (mlirDialectIsNull(dialect)) {
+    throw SetPyError(attrError ? PyExc_AttributeError : PyExc_IndexError,
+                     llvm::Twine("Dialect '") + key + "' not found");
+  }
+  return dialect;
+}
+
 //------------------------------------------------------------------------------
 // PyModule
 //------------------------------------------------------------------------------
 
 PyModule::PyModule(PyMlirContextRef contextRef, MlirModule module)
     : BaseContextObject(std::move(contextRef)), module(module) {}
 
 PyModule::~PyModule() {

 };
 
 } // namespace
 
 //------------------------------------------------------------------------------
 // Populates the pybind11 IR submodule.
 //------------------------------------------------------------------------------
 
-void mlir::python::populateIRSubmodule(py::module &m) {
+void mlir::python::populateIRSubmodule(py::module m) {
+  //----------------------------------------------------------------------------
+  // Module local extension management.
+  // This is complicated because it is python state carried at the module
+  // level, effectively as statics. In order to ensure proper cleanup and
+  // sharing (in the case of defining multiple modules), this is managed
+  // via python constructs, making the code somewhat more complicated.
+  //----------------------------------------------------------------------------
+
+  // The list of module names to search for python extensions to standard
+  // classes. Each element in the list is the name of a parent module,
+  // containing sub-modules with well-defined names relating to various
+  // extensions. Modules are search in order for matches.
+  m.attr("extension_search_list") = py::list();
+
+  // Extensions modules loaded against this module, keyed by extension name.
+  // When an extension is not in the dict, it is to be referenced by searching
+  // for submodules with the given name in the extension_search_list. Once
+  // an extension has been loaded, it will be remembered, even if the search
+  // list changes.
+  m.attr("_extensions") = py::dict();
+
+  // Loads an extension module by extension name. This is done by searching
+  // the extension_search_list and constructing a derived module as:
+  //   <searchModule>.<extensionName>
+  // Returns None on failure to find.
+  auto loadExtension = [m](py::str extensionName) -> py::module {
+    py::dict extensions = m.attr("_extensions");
+    py::module extModule = extensions.attr("get")(extensionName);
+    if (!extModule.is_none())
+      return extModule;
+
+    auto searchList = m.attr("extension_search_list");
+    for (auto it : searchList) {
+      std::string parentModule = py::str(it);
+      parentModule.append(".");
+      parentModule.append(py::str(extensionName));
+      try {
+        extModule = py::module::import(parentModule.c_str());
+      } catch (py::error_already_set &e) {
+        if (e.matches(PyExc_ImportError)) {
+          continue;
+        }
+        throw;
+      }
+      extensions[extensionName] = extModule;
+      return extModule;
+    }
+
+    return py::none();
+  };
+
+  auto createDialectExtension = [=](const std::string &dialectName,
+                                    py::object descriptor) -> py::object {
+    std::string extensionName = dialectName;
+    extensionName.append("_dialect");
+    py::module extModule = loadExtension(extensionName);
+    if (extModule.is_none()) {
+      // Return a non extension base version.
+      return py::cast(PyDialect{std::move(descriptor)});
+    } else {
+      // Instantiate the "Dialect" class in the extension.
+      return extModule.attr("Dialect")(std::move(descriptor));
+    }
+  };
+
+  //----------------------------------------------------------------------------
   // Mapping of MlirContext
+  //----------------------------------------------------------------------------
   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_readonly(
+          "dialect",
+          [](PyMlirContext &self) { return PyDialects(self.getRef()); },
+          "Gets a container for accessing dialects by name")
+      .def_property_readonly(
+          "d", [](PyMlirContext &self) { return PyDialects(self.getRef()); },
+          "Alias for 'dialect'")
+      .def(
+          "get_dialect_descriptor",
+          [=](PyMlirContext &self, std::string &name) {
+            MlirDialect dialect = mlirContextGetOrLoadDialect(
+                self.get(), {name.data(), name.size()});
+            if (mlirDialectIsNull(dialect)) {
+              throw SetPyError(PyExc_ValueError,
+                               llvm::Twine("Dialect '") + name + "' not found");
+            }
+            return PyDialectDescriptor(self.getRef(), dialect);
+          },
+          "Gets or loads a dialect by name, returning its descriptor object")
       .def_property(
           "allow_unregistered_dialects",
           [](PyMlirContext &self) -> bool {
             return mlirContextGetAllowUnregisteredDialects(self.get());
           },
           [](PyMlirContext &self, bool value) {
             mlirContextSetAllowUnregisteredDialects(self.get(), value);
           })
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"));
 
+  //----------------------------------------------------------------------------
+  // Mapping of PyDialectDescriptor
+  //----------------------------------------------------------------------------
+  py::class_<PyDialectDescriptor>(m, "DialectDescriptor")
+      .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")
+      .def("__getitem__",
+           [=](PyDialects &self, std::string keyName) {
+             MlirDialect dialect =
+                 self.getDialectForKey(keyName, /*attrError=*/false);
+             py::object descriptor =
+                 py::cast(PyDialectDescriptor{self.getContext(), dialect});
+             return createDialectExtension(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 createDialectExtension(attrName, std::move(descriptor));
+      });
+
+  //----------------------------------------------------------------------------
+  // Mapping of PyDialect
+  //----------------------------------------------------------------------------
+  py::class_<PyDialect>(m, "Dialect")
+      .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")
       .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(
py::class_<PyModule>(m, "Module")
             MlirOperation operation = mlirModuleGetOperation(self.get());
             PyPrintAccumulator printAccum;
             mlirOperationPrint(operation, printAccum.getCallback(),
                                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()); })
py::class_<PyOperation>(m, "Operation")
            // Careful: Lots of arguments must match up with get_asm method.
            py::arg("binary") = false,
            py::arg("large_elements_limit") = py::none(),
            py::arg("enable_debug_info") = false,
            py::arg("pretty_debug_info") = false,
            py::arg("print_generic_op_form") = false,
            py::arg("use_local_scope") = false, kOperationGetAsmDocstring);
 
+  //----------------------------------------------------------------------------
   // Mapping of PyRegion.
+  //----------------------------------------------------------------------------
   py::class_<PyRegion>(m, "Region")
       .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, py::object &other) {
         try {
           PyRegion *otherRegion = other.cast<PyRegion *>();
           return self.get().ptr == otherRegion->get().ptr;
         } catch (std::exception &e) {
           return false;
         }
       });
 
+  //----------------------------------------------------------------------------
   // Mapping of PyBlock.
+  //----------------------------------------------------------------------------
   py::class_<PyBlock>(m, "Block")
       .def_property_readonly(
           "arguments",
           [](PyBlock &self) {
             return PyBlockArgumentList(self.getParentOperation(), self.get());
           },
           "Returns a list of block arguments.")
       .def_property_readonly(
py::class_<PyBlock>(m, "Block")
             self.checkValid();
             PyPrintAccumulator printAccum;
             mlirBlockPrint(self.get(), printAccum.getCallback(),
                            printAccum.getUserData());
             return printAccum.join();
           },
           "Returns the assembly form of the block.");
 
+  //----------------------------------------------------------------------------
   // Mapping of PyAttribute.
+  //----------------------------------------------------------------------------
   py::class_<PyAttribute>(m, "Attribute")
       .def_property_readonly(
           "context",
           [](PyAttribute &self) { return self.getContext().getObject(); },
           "Context that owns the Attribute")
       .def_property_readonly("type",
                              [](PyAttribute &self) {
                                return PyType(self.getContext()->getRef(),
py::class_<PyAttribute>(m, "Attribute")
         PyPrintAccumulator printAccum;
         printAccum.parts.append("Attribute(");
         mlirAttributePrint(self.attr, printAccum.getCallback(),
                            printAccum.getUserData());
         printAccum.parts.append(")");
         return printAccum.join();
       });
 
+  //----------------------------------------------------------------------------
+  // Mapping of PyNamedAttribute
+  //----------------------------------------------------------------------------
   py::class_<PyNamedAttribute>(m, "NamedAttribute")
       .def("__repr__",
            [](PyNamedAttribute &self) {
              PyPrintAccumulator printAccum;
              printAccum.parts.append("NamedAttribute(");
              printAccum.parts.append(self.namedAttr.name);
              printAccum.parts.append("=");
              mlirAttributePrint(self.namedAttr.attribute,
void mlir::python::populateIRSubmodule(py::module m) {
 
   // Standard attribute bindings.
   PyFloatAttribute::bind(m);
   PyIntegerAttribute::bind(m);
   PyBoolAttribute::bind(m);
   PyStringAttribute::bind(m);
   PyDenseElementsAttribute::bind(m);
 
+  //----------------------------------------------------------------------------
   // Mapping of PyType.
+  //----------------------------------------------------------------------------
   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>();
void mlir::python::populateIRSubmodule(py::module m) {
   PyVectorType::bind(m);
   PyRankedTensorType::bind(m);
   PyUnrankedTensorType::bind(m);
   PyMemRefType::bind(m);
   PyUnrankedMemRefType::bind(m);
   PyTupleType::bind(m);
   PyFunctionType::bind(m);
 
+  //----------------------------------------------------------------------------
   // Mapping of Value.
+  //----------------------------------------------------------------------------
   py::class_<PyValue>(m, "Value")
       .def_property_readonly(
           "context",
           [](PyValue &self) { return self.getParentOperation()->getContext(); },
           "Context in which the value lives.")
       .def(
           "dump", [](PyValue &self) { mlirValueDump(self.get()); },
           kDumpDocstring)

mlir/lib/Bindings/Python/mlir/__init__.py

 #  Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 #  See https://llvm.org/LICENSE.txt for license information.
 #  SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
 # Note that the only function of this module is currently to load the
 # native module and re-export its symbols. In the future, this file is
 # reserved as a trampoline to handle environment specific loading needs
 # and arbitrate any one-time initialization needed in various shared-library
 # scenarios.
 
 from _mlir import *
+
+# Add our dialect package to the module search list.
+ir.extension_search_list.append("mlir.extensions")

mlir/lib/Bindings/Python/mlir/extensions/std_dialect.py

+#  Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+#  See https://llvm.org/LICENSE.txt for license information.
+#  SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+from mlir import ir
+
+class Dialect(ir.Dialect):
+  pass
+
+
+# Example custom op:
+#   - The class decorator does some meta-programming to clean up its API and
+#     sets it as an attr on the static Dialect.
+#   - Can be accessed with "ctx.dialect.std.AddFOp"
+# @Dialect.register_op
+# class AddFOp(ir.Op):
+#   @staticmethod
+#   def create(ctx, loc, lhs, rhs):
+#     return ctx.create_operation(
+#         "std.addf", loc, operands=[lhs, rhs], results=[lhs.type])
+

mlir/test/Bindings/Python/dialects.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: testDialectDescriptor
+def testDialectDescriptor():
+  ctx = mlir.ir.Context()
+  d = ctx.get_dialect_descriptor("std")
+  # CHECK: <DialectDescriptor std>
+  print(d)
+  # CHECK: std
+  print(d.namespace)
+  try:
+    _ = ctx.get_dialect_descriptor("not_existing")
+  except ValueError:
+    pass
+  else:
+    assert False, "Expected exception"
+
+run(testDialectDescriptor)
+
+
+# CHECK-LABEL: TEST: testUserDialects
+def testUserDialects():
+  ctx = mlir.ir.Context()
+  # Access using attribute.
+  d = ctx.dialect.std
+  # CHECK: <Dialect std (class mlir.extensions.std_dialect.Dialect)>
+  print(d)
+  try:
+    _ = ctx.dialect.not_existing
+  except AttributeError:
+    pass
+  else:
+    assert False, "Expected exception"
+
+  # Access using index.
+  d = ctx.dialect["std"]
+  # CHECK: <Dialect std (class mlir.extensions.std_dialect.Dialect)>
+  print(d)
+  try:
+    _ = ctx.dialect["not_existing"]
+  except IndexError:
+    pass
+  else:
+    assert False, "Expected exception"
+
+  # Using the 'd' alias.
+  d = ctx.d["std"]
+  # CHECK: <Dialect std (class mlir.extensions.std_dialect.Dialect)>
+  print(d)
+
+
+run(testUserDialects)