[mlir][Python] Initial Affine Map Python Bindings.

[zhanghb97]

Summary

• Add PyAffineMap to wrap around MlirAffineMap.
• Add mlirPythonAffineMapToCapsule and mlirPythonCapsuleToAffineMap to interoperate with python capsule.
• Add and test some simple bindings of PyAffineMap.

[zhanghb97]

Rebase.

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/AffineMap.h"
 #include "mlir-c/IR.h"
 #include "mlir-c/Pass.h"
 
+#define MLIR_PYTHON_CAPSULE_AFFINE_MAP "mlir.ir.AffineMap._CAPIPtr"
 #define MLIR_PYTHON_CAPSULE_ATTRIBUTE "mlir.ir.Attribute._CAPIPtr"
 #define MLIR_PYTHON_CAPSULE_CONTEXT "mlir.ir.Context._CAPIPtr"
 #define MLIR_PYTHON_CAPSULE_LOCATION "mlir.ir.Location._CAPIPtr"
 #define MLIR_PYTHON_CAPSULE_MODULE "mlir.ir.Module._CAPIPtr"
 #define MLIR_PYTHON_CAPSULE_OPERATION "mlir.ir.Operation._CAPIPtr"
 #define MLIR_PYTHON_CAPSULE_TYPE "mlir.ir.Type._CAPIPtr"
 #define MLIR_PYTHON_CAPSULE_PASS_MANAGER "mlir.passmanager.PassManager._CAPIPtr"
 

  * a null type is returned (as checked via mlirTypeIsNull). In such a
  * case, the Python APIs will have already set an error. */
 static inline MlirType mlirPythonCapsuleToType(PyObject *capsule) {
   void *ptr = PyCapsule_GetPointer(capsule, MLIR_PYTHON_CAPSULE_TYPE);
   MlirType type = {ptr};
   return type;
 }
 
+/** Creates a capsule object encapsulating the raw C-API MlirAffineMap.
+ * The returned capsule does not extend or affect ownership of any Python
+ * objects that reference the type in any way.
+ */
+static inline PyObject *mlirPythonAffineMapToCapsule(MlirAffineMap affineMap) {
+  return PyCapsule_New(MLIR_PYTHON_GET_WRAPPED_POINTER(affineMap),
+                       MLIR_PYTHON_CAPSULE_AFFINE_MAP, NULL);
+}
+
+/** Extracts an MlirAffineMap from a capsule as produced from
+ * mlirPythonAffineMapToCapsule. If the capsule is not of the right type, then
+ * a null type is returned (as checked via mlirAffineMapIsNull). In such a
+ * case, the Python APIs will have already set an error. */
+static inline MlirAffineMap mlirPythonCapsuleToAffineMap(PyObject *capsule) {
+  void *ptr = PyCapsule_GetPointer(capsule, MLIR_PYTHON_CAPSULE_AFFINE_MAP);
+  MlirAffineMap affineMap = {ptr};
+  return affineMap;
+}
+
 #ifdef __cplusplus
 }
 #endif
 
 #endif // MLIR_C_BINDINGS_PYTHON_INTEROP_H

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 <vector>
 
 #include "PybindUtils.h"
 
+#include "mlir-c/AffineMap.h"
 #include "mlir-c/IR.h"
 #include "llvm/ADT/DenseMap.h"
 
 namespace mlir {
 namespace python {
 
 class PyBlock;
 class PyInsertionPoint;
public:
 
   void checkValid() { return parentOperation->checkValid(); }
 
 private:
   PyOperationRef parentOperation;
   MlirValue value;
 };
 
+class PyAffineMap : public BaseContextObject {
+public:
+  PyAffineMap(PyMlirContextRef contextRef, MlirAffineMap affineMap)
+      : BaseContextObject(std::move(contextRef)), affineMap(affineMap) {}
+  bool operator==(const PyAffineMap &other);
+  operator MlirAffineMap() const { return affineMap; }
+  MlirAffineMap get() const { return affineMap; }
+
+  /// Gets a capsule wrapping the void* within the MlirAffineMap.
+  pybind11::object getCapsule();
+
+  /// Creates a PyAffineMap from the MlirAffineMap wrapped by a capsule.
+  /// Note that PyAffineMap instances are uniqued, so the returned object
+  /// may be a pre-existing object. Ownership of the underlying MlirAffineMap
+  /// is taken by calling this function.
+  static PyAffineMap createFromCapsule(pybind11::object capsule);
+
+private:
+  MlirAffineMap affineMap;
+};
+
 void populateIRSubmodule(pybind11::module &m);
 
 } // namespace python
 } // namespace mlir
 
 namespace pybind11 {
 namespace detail {
 

mlir/lib/Bindings/Python/IRModules.cpp

c.def_property_readonly(
         },
         "Returns the list of result types in the FunctionType.");
   }
 };
 
 } // namespace
 
 //------------------------------------------------------------------------------
+// PyAffineMap.
+//------------------------------------------------------------------------------
+
+bool PyAffineMap::operator==(const PyAffineMap &other) {
+  return mlirAffineMapEqual(affineMap, other.affineMap);
+}
+
+py::object PyAffineMap::getCapsule() {
+  return py::reinterpret_steal<py::object>(mlirPythonAffineMapToCapsule(*this));
+}
+
+PyAffineMap PyAffineMap::createFromCapsule(py::object capsule) {
+  MlirAffineMap rawAffineMap = mlirPythonCapsuleToAffineMap(capsule.ptr());
+  if (mlirAffineMapIsNull(rawAffineMap))
+    throw py::error_already_set();
+  return PyAffineMap(
+      PyMlirContext::forContext(mlirAffineMapGetContext(rawAffineMap)),
+      rawAffineMap);
+}
+
+//------------------------------------------------------------------------------
 // Populates the pybind11 IR submodule.
 //------------------------------------------------------------------------------
 
 void mlir::python::populateIRSubmodule(py::module &m) {
   //----------------------------------------------------------------------------
   // Mapping of MlirContext
   //----------------------------------------------------------------------------
   py::class_<PyMlirContext>(m, "Context")
void mlir::python::populateIRSubmodule(py::module &m) {
   PyBlockList::bind(m);
   PyOperationIterator::bind(m);
   PyOperationList::bind(m);
   PyOpAttributeMap::bind(m);
   PyOpOperandList::bind(m);
   PyOpResultList::bind(m);
   PyRegionIterator::bind(m);
   PyRegionList::bind(m);
+
+  //----------------------------------------------------------------------------
+  // Mapping of PyAffineMap.
+  //----------------------------------------------------------------------------
+  py::class_<PyAffineMap>(m, "AffineMap")
+      .def_property_readonly(MLIR_PYTHON_CAPI_PTR_ATTR,
+                             &PyAffineMap::getCapsule)
+      .def(MLIR_PYTHON_CAPI_FACTORY_ATTR, &PyAffineMap::createFromCapsule)
+      .def_static(
+          "get_empty",
+          [](DefaultingPyMlirContext context) {
+            MlirAffineMap affineMap = mlirAffineMapEmptyGet(context->get());
+            return PyAffineMap(context->getRef(), affineMap);
+          },
+          py::arg("context") = py::none(), "Gets an empty affine map.")
+      .def_property_readonly(
+          "context",
+          [](PyAffineMap &self) { return self.getContext().getObject(); },
+          "Context that owns the Affine Map")
+      .def("__eq__",
+           [](PyAffineMap &self, PyAffineMap &other) { return self == other; })
+      .def("__eq__", [](PyAffineMap &self, py::object &other) { return false; })
+      .def(
+          "dump", [](PyAffineMap &self) { mlirAffineMapDump(self); },
+          kDumpDocstring);
 }

mlir/test/Bindings/Python/ir_affine_map.py

+# RUN: %PYTHON %s | FileCheck %s
+
+import gc
+from mlir.ir import *
+
+def run(f):
+  print("\nTEST:", f.__name__)
+  f()
+  gc.collect()
+  assert Context._get_live_count() == 0
+
+
+# CHECK-LABEL: TEST: testAffineMapCapsule
+def testAffineMapCapsule():
+  with Context() as ctx:
+    am1 = AffineMap.get_empty(ctx)
+  # CHECK: mlir.ir.AffineMap._CAPIPtr
+  affine_map_capsule = am1._CAPIPtr
+  print(affine_map_capsule)
+  am2 = AffineMap._CAPICreate(affine_map_capsule)
+  assert am2 == am1
+  assert am2.context is ctx
+
+run(testAffineMapCapsule)