[mlir] Add Python bindings for IntegerSet

[ftynse]

Summary

This follows up on the introduction of C API for the same object and is similar
to AffineExpr and AffineMap.

[stellaraccident]

Unrelated: the IRModules.cpp file is getting quite large and slow to compile. Any objection to splitting it into a few pieces?

mlir/lib/Bindings/Python/IRModules.cpp

4367

I believe the eq method should return False on mismatched types, not throw. There is a simple overload that does it: see other examples in this file.

4415

Nit: shouldn't pt args be snake case? (Here and below)

mlir/test/Bindings/Python/ir_integer_set.py

37

Check intended? (Here and below)

[ftynse]

Unrelated: the IRModules.cpp file is getting quite large and slow to compile. Any objection to splitting it into a few pieces?

I was going to propose the same thing.

[ftynse]

Address review

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/AffineExpr.h"
 #include "mlir-c/AffineMap.h"
 #include "mlir-c/IR.h"
+#include "mlir-c/IntegerSet.h"
 #include "mlir-c/Pass.h"
 
 #define MLIR_PYTHON_CAPSULE_AFFINE_EXPR "mlir.ir.AffineExpr._CAPIPtr"
 #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_INTEGER_SET "mlir.ir.IntegerSet._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"
 
 /** 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

  * 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;
 }
 
+/** Creates a capsule object encapsulating the raw C-API MlirIntegerSet.
+ * The returned capsule does not extend or affect ownership of any Python
+ * objects that reference the set in any way. */
+static inline PyObject *
+mlirPythonIntegerSetToCapsule(MlirIntegerSet integerSet) {
+  return PyCapsule_New(MLIR_PYTHON_GET_WRAPPED_POINTER(integerSet),
+                       MLIR_PYTHON_CAPSULE_INTEGER_SET, NULL);
+}
+
+/** Extracts an MlirIntegerSet from a capsule as produced from
+ * mlirPythonIntegerSetToCapsule. If the capsule is not of the right type, then
+ * a null set is returned (as checked via mlirIntegerSetIsNull). In such a
+ * case, the Python APIs will have already set an error. */
+static inline MlirIntegerSet mlirPythonCapsuleToIntegerSet(PyObject *capsule) {
+  void *ptr = PyCapsule_GetPointer(capsule, MLIR_PYTHON_CAPSULE_INTEGER_SET);
+  MlirIntegerSet integerSet = {ptr};
+  return integerSet;
+}
+
 #ifdef __cplusplus
 }
 #endif
 
 #endif // MLIR_C_BINDINGS_PYTHON_INTEROP_H

mlir/lib/Bindings/Python/IRModules.h

 
 #include <vector>
 
 #include "PybindUtils.h"
 
 #include "mlir-c/AffineExpr.h"
 #include "mlir-c/AffineMap.h"
 #include "mlir-c/IR.h"
+#include "mlir-c/IntegerSet.h"
 #include "llvm/ADT/DenseMap.h"
 
 namespace mlir {
 namespace python {
 
 class PyBlock;
 class PyInsertionPoint;
 class PyLocation;
public:
   /// 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;
 };
 
+class PyIntegerSet : public BaseContextObject {
+public:
+  PyIntegerSet(PyMlirContextRef contextRef, MlirIntegerSet integerSet)
+      : BaseContextObject(std::move(contextRef)), integerSet(integerSet) {}
+  bool operator==(const PyIntegerSet &other);
+  operator MlirIntegerSet() const { return integerSet; }
+  MlirIntegerSet get() const { return integerSet; }
+
+  /// Gets a capsule wrapping the void* within the MlirIntegerSet.
+  pybind11::object getCapsule();
+
+  /// Creates a PyIntegerSet from the MlirAffineMap wrapped by a capsule.
+  /// Note that PyIntegerSet instances may be uniqued, so the returned object
+  /// may be a pre-existing object. Integer sets are owned by the context.
+  static PyIntegerSet createFromCapsule(pybind11::object capsule);
+
+private:
+  MlirIntegerSet integerSet;
+};
+
 void populateIRSubmodule(pybind11::module &m);
 
 } // namespace python
 } // namespace mlir
 
 namespace pybind11 {
 namespace detail {
 

mlir/lib/Bindings/Python/IRModules.cpp

 
 #include "Globals.h"
 #include "PybindUtils.h"
 
 #include "mlir-c/AffineMap.h"
 #include "mlir-c/Bindings/Python/Interop.h"
 #include "mlir-c/BuiltinAttributes.h"
 #include "mlir-c/BuiltinTypes.h"
+#include "mlir-c/IntegerSet.h"
 #include "mlir-c/Registration.h"
 #include "llvm/ADT/SmallVector.h"
 #include <pybind11/stl.h>
 
 namespace py = pybind11;
 using namespace mlir;
 using namespace mlir::python;
 
PyAffineMap PyAffineMap::createFromCapsule(py::object capsule) {
   if (mlirAffineMapIsNull(rawAffineMap))
     throw py::error_already_set();
   return PyAffineMap(
       PyMlirContext::forContext(mlirAffineMapGetContext(rawAffineMap)),
       rawAffineMap);
 }
 
 //------------------------------------------------------------------------------
+// PyIntegerSet and utilities.
+//------------------------------------------------------------------------------
+
+class PyIntegerSetConstraint {
+public:
+  PyIntegerSetConstraint(PyIntegerSet set, intptr_t pos) : set(set), pos(pos) {}
+
+  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")
+        .def_property_readonly("expr", &PyIntegerSetConstraint::getExpr)
+        .def_property_readonly("is_eq", &PyIntegerSetConstraint::isEq);
+  }
+
+private:
+  PyIntegerSet set;
+  intptr_t pos;
+};
+
+class PyIntegerSetConstraintList
+    : public Sliceable<PyIntegerSetConstraintList, PyIntegerSetConstraint> {
+public:
+  static constexpr const char *pyClassName = "IntegerSetConstraintList";
+
+  PyIntegerSetConstraintList(PyIntegerSet set, intptr_t startIndex = 0,
+                             intptr_t length = -1, intptr_t step = 1)
+      : Sliceable(startIndex,
+                  length == -1 ? mlirIntegerSetGetNumConstraints(set) : length,
+                  step),
+        set(set) {}
+
+  intptr_t getNumElements() { return mlirIntegerSetGetNumConstraints(set); }
+
+  PyIntegerSetConstraint getElement(intptr_t pos) {
+    return PyIntegerSetConstraint(set, pos);
+  }
+
+  PyIntegerSetConstraintList slice(intptr_t startIndex, intptr_t length,
+                                   intptr_t step) {
+    return PyIntegerSetConstraintList(set, startIndex, length, step);
+  }
+
+private:
+  PyIntegerSet set;
+};
+
+bool PyIntegerSet::operator==(const PyIntegerSet &other) {
+  return mlirIntegerSetEqual(integerSet, other.integerSet);
+}
+
+py::object PyIntegerSet::getCapsule() {
+  return py::reinterpret_steal<py::object>(
+      mlirPythonIntegerSetToCapsule(*this));
+}
+
+PyIntegerSet PyIntegerSet::createFromCapsule(py::object capsule) {
+  MlirIntegerSet rawIntegerSet = mlirPythonCapsuleToIntegerSet(capsule.ptr());
+  if (mlirIntegerSetIsNull(rawIntegerSet))
+    throw py::error_already_set();
+  return PyIntegerSet(
+      PyMlirContext::forContext(mlirIntegerSetGetContext(rawIntegerSet)),
+      rawIntegerSet);
+}
+
+/// Attempts to populate `result` with the content of `list` casted to the
+/// appropriate type (Python and C types are provided as template arguments).
+/// Throws errors in case of failure, using "action" to describe what the caller
+/// was attempting to do.
+template <typename PyType, typename CType>
+static void pyListToVector(py::list list, llvm::SmallVectorImpl<CType> &result,
+                           StringRef action) {
+  result.reserve(py::len(list));
+  for (py::handle item : list) {
+    try {
+      result.push_back(item.cast<PyType>());
+    } catch (py::cast_error &err) {
+      std::string msg = (llvm::Twine("Invalid expression when ") + action +
+                         " (" + err.what() + ")")
+                            .str();
+      throw py::cast_error(msg);
+    } catch (py::reference_cast_error &err) {
+      std::string msg = (llvm::Twine("Invalid expression (None?) when ") +
+                         action + " (" + err.what() + ")")
+                            .str();
+      throw py::cast_error(msg);
+    }
+  }
+}
+
+//------------------------------------------------------------------------------
 // Populates the pybind11 IR submodule.
 //------------------------------------------------------------------------------
 
 void mlir::python::populateIRSubmodule(py::module &m) {
   //----------------------------------------------------------------------------
   // Mapping of MlirContext
   //----------------------------------------------------------------------------
   py::class_<PyMlirContext>(m, "Context")
py::class_<PyAffineMap>(m, "AffineMap")
       .def(
           "dump", [](PyAffineMap &self) { mlirAffineMapDump(self); },
           kDumpDocstring)
       .def_static(
           "get",
           [](intptr_t dimCount, intptr_t symbolCount, py::list exprs,
              DefaultingPyMlirContext context) {
             SmallVector<MlirAffineExpr> affineExprs;
-            affineExprs.reserve(py::len(exprs));
-            for (py::handle expr : exprs) {
-              try {
-                affineExprs.push_back(expr.cast<PyAffineExpr>());
-              } catch (py::cast_error &err) {
-                std::string msg =
-                    std::string("Invalid expression when attempting to create "
-                                "an AffineMap (") +
-                    err.what() + ")";
-                throw py::cast_error(msg);
-              } catch (py::reference_cast_error &err) {
-                std::string msg =
-                    std::string("Invalid expression (None?) when attempting to "
-                                "create an AffineMap (") +
-                    err.what() + ")";
-                throw py::cast_error(msg);
-              }
-            }
+            pyListToVector<PyAffineExpr, MlirAffineExpr>(
+                exprs, affineExprs, "attempting to create an AffineMap");
             MlirAffineMap map =
                 mlirAffineMapGet(context->get(), dimCount, symbolCount,
                                  affineExprs.size(), affineExprs.data());
             return PyAffineMap(context->getRef(), map);
           },
           py::arg("dim_count"), py::arg("symbol_count"), py::arg("exprs"),
           py::arg("context") = py::none(),
           "Gets a map with the given expressions as results.")
py::class_<PyAffineMap>(m, "AffineMap")
           [](PyAffineMap &self) { return mlirAffineMapGetNumInputs(self); })
       .def_property_readonly(
           "n_symbols",
           [](PyAffineMap &self) { return mlirAffineMapGetNumSymbols(self); })
       .def_property_readonly("results", [](PyAffineMap &self) {
         return PyAffineMapExprList(self);
       });
   PyAffineMapExprList::bind(m);
+
+  //----------------------------------------------------------------------------
+  // Mapping of PyIntegerSet.
+  //----------------------------------------------------------------------------
+  py::class_<PyIntegerSet>(m, "IntegerSet")
+      .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) {
+             PyPrintAccumulator printAccum;
+             mlirIntegerSetPrint(self, printAccum.getCallback(),
+                                 printAccum.getUserData());
+             return printAccum.join();
+           })
+      .def("__repr__",
+           [](PyIntegerSet &self) {
+             PyPrintAccumulator printAccum;
+             printAccum.parts.append("IntegerSet(");
+             mlirIntegerSetPrint(self, printAccum.getCallback(),
+                                 printAccum.getUserData());
+             printAccum.parts.append(")");
+             return printAccum.join();
+           })
+      .def_property_readonly(
+          "context",
+          [](PyIntegerSet &self) { return self.getContext().getObject(); })
+      .def(
+          "dump", [](PyIntegerSet &self) { mlirIntegerSetDump(self); },
+          kDumpDocstring)
+      .def_static(
+          "get",
+          [](intptr_t numDims, intptr_t numSymbols, py::list exprs,
+             std::vector<bool> eqFlags, DefaultingPyMlirContext context) {
+            if (exprs.size() != eqFlags.size())
+              throw py::value_error(
+                  "Expected the number of constraints to match "
+                  "that of equality flags");
+            if (exprs.empty())
+              throw py::value_error("Expected non-empty list of constraints");
+
+            // Copy over to a SmallVector because std::vector has a
+            // specialization for booleans that packs data and does not
+            // expose a `bool *`.
+            SmallVector<bool, 8> flags(eqFlags.begin(), eqFlags.end());
+
+            SmallVector<MlirAffineExpr> affineExprs;
+            pyListToVector<PyAffineExpr>(exprs, affineExprs,
+                                         "attempting to create an IntegerSet");
+            MlirIntegerSet set = mlirIntegerSetGet(
+                context->get(), numDims, numSymbols, exprs.size(),
+                affineExprs.data(), flags.data());
+            return PyIntegerSet(context->getRef(), set);
+          },
+          py::arg("num_dims"), py::arg("num_symbols"), py::arg("exprs"),
+          py::arg("eq_flags"), py::arg("context") = py::none())
+      .def_static(
+          "get_empty",
+          [](intptr_t numDims, intptr_t numSymbols,
+             DefaultingPyMlirContext context) {
+            MlirIntegerSet set =
+                mlirIntegerSetEmptyGet(context->get(), numDims, numSymbols);
+            return PyIntegerSet(context->getRef(), set);
+          },
+          py::arg("num_dims"), py::arg("num_symbols"),
+          py::arg("context") = py::none())
+      .def("get_replaced",
+           [](PyIntegerSet &self, py::list dimExprs, py::list symbolExprs,
+              intptr_t numResultDims, intptr_t numResultSymbols) {
+             if (static_cast<intptr_t>(dimExprs.size()) !=
+                 mlirIntegerSetGetNumDims(self))
+               throw py::value_error(
+                   "Expected the number of dimension replacement expressions "
+                   "to match that of dimensions");
+             if (static_cast<intptr_t>(symbolExprs.size()) !=
+                 mlirIntegerSetGetNumSymbols(self))
+               throw py::value_error(
+                   "Expected the number of symbol replacement expressions "
+                   "to match that of symbols");
+
+             SmallVector<MlirAffineExpr> dimAffineExprs, symbolAffineExprs;
+             pyListToVector<PyAffineExpr>(
+                 dimExprs, dimAffineExprs,
+                 "attempting to create an IntegerSet by replacing dimensions");
+             pyListToVector<PyAffineExpr>(
+                 symbolExprs, symbolAffineExprs,
+                 "attempting to create an IntegerSet by replacing symbols");
+             MlirIntegerSet set = mlirIntegerSetReplaceGet(
+                 self, dimAffineExprs.data(), symbolAffineExprs.data(),
+                 numResultDims, numResultSymbols);
+             return PyIntegerSet(self.getContext(), set);
+           })
+      .def_property_readonly("is_canonical_empty",
+                             [](PyIntegerSet &self) {
+                               return mlirIntegerSetIsCanonicalEmpty(self);
+                             })
+      .def_property_readonly(
+          "n_dims",
+          [](PyIntegerSet &self) { return mlirIntegerSetGetNumDims(self); })
+      .def_property_readonly(
+          "n_symbols",
+          [](PyIntegerSet &self) { return mlirIntegerSetGetNumSymbols(self); })
+      .def_property_readonly(
+          "n_inputs",
+          [](PyIntegerSet &self) { return mlirIntegerSetGetNumInputs(self); })
+      .def_property_readonly("n_equalities",
+                             [](PyIntegerSet &self) {
+                               return mlirIntegerSetGetNumEqualities(self);
+                             })
+      .def_property_readonly("n_inequalities",
+                             [](PyIntegerSet &self) {
+                               return mlirIntegerSetGetNumInequalities(self);
+                             })
+      .def_property_readonly("constraints", [](PyIntegerSet &self) {
+        return PyIntegerSetConstraintList(self);
+      });
+  PyIntegerSetConstraint::bind(m);
+  PyIntegerSetConstraintList::bind(m);
 }

mlir/test/Bindings/Python/ir_integer_set.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: testIntegerSetCapsule
+def testIntegerSetCapsule():
+  with Context() as ctx:
+    is1 = IntegerSet.get_empty(1, 1, ctx)
+  capsule = is1._CAPIPtr
+  # CHECK: mlir.ir.IntegerSet._CAPIPtr
+  print(capsule)
+  is2 = IntegerSet._CAPICreate(capsule)
+  assert is1 == is2
+  assert is2.context is ctx
+
+run(testIntegerSetCapsule)
+
+
+# CHECK-LABEL: TEST: testIntegerSetGet
+def testIntegerSetGet():
+  with Context():
+    d0 = AffineDimExpr.get(0)
+    d1 = AffineDimExpr.get(1)
+    s0 = AffineSymbolExpr.get(0)
+    c42 = AffineConstantExpr.get(42)
+
+    # CHECK: (d0, d1)[s0] : (d0 - d1 == 0, s0 - 42 >= 0)
+    set0 = IntegerSet.get(2, 1, [d0 - d1, s0 - c42], [True, False])
+    print(set0)
+
+    # CHECK: (d0)[s0] : (1 == 0)
+    set1 = IntegerSet.get_empty(1, 1)
+    print(set1)
+
+    # CHECK: (d0)[s0, s1] : (d0 - s1 == 0, s0 - 42 >= 0)
+    set2 = set0.get_replaced([d0, AffineSymbolExpr.get(1)], [s0], 1, 2)
+    print(set2)
+
+    try:
+      IntegerSet.get(2, 1, [], [])
+    except ValueError as e:
+      # CHECK: Expected non-empty list of constraints
+      print(e)
+
+    try:
+      IntegerSet.get(2, 1, [d0 - d1], [True, False])
+    except ValueError as e:
+      # CHECK: Expected the number of constraints to match that of equality flags
+      print(e)
+
+    try:
+      IntegerSet.get(2, 1, [0], [True])
+    except RuntimeError as e:
+      # CHECK: Invalid expression when attempting to create an IntegerSet
+      print(e)
+
+    try:
+      IntegerSet.get(2, 1, [None], [True])
+    except RuntimeError as e:
+      # CHECK: Invalid expression (None?) when attempting to create an IntegerSet
+      print(e)
+
+    try:
+      set0.get_replaced([d0], [s0], 1, 1)
+    except ValueError as e:
+      # CHECK: Expected the number of dimension replacement expressions to match that of dimensions
+      print(e)
+
+    try:
+      set0.get_replaced([d0, d1], [s0, s0], 1, 1)
+    except ValueError as e:
+      # CHECK: Expected the number of symbol replacement expressions to match that of symbols
+      print(e)
+
+    try:
+      set0.get_replaced([d0, 1], [s0], 1, 1)
+    except RuntimeError as e:
+      # CHECK: Invalid expression when attempting to create an IntegerSet by replacing dimensions
+      print(e)
+
+    try:
+      set0.get_replaced([d0, d1], [None], 1, 1)
+    except RuntimeError as e:
+      # CHECK: Invalid expression (None?) when attempting to create an IntegerSet by replacing symbols
+      print(e)
+
+run(testIntegerSetGet)
+
+
+# CHECK-LABEL: TEST: testIntegerSetProperties
+def testIntegerSetProperties():
+  with Context():
+    d0 = AffineDimExpr.get(0)
+    d1 = AffineDimExpr.get(1)
+    s0 = AffineSymbolExpr.get(0)
+    c42 = AffineConstantExpr.get(42)
+
+    set0 = IntegerSet.get(2, 1, [d0 - d1, s0 - c42, s0 - d0], [True, False, False])
+    # CHECK: 2
+    print(set0.n_dims)
+    # CHECK: 1
+    print(set0.n_symbols)
+    # CHECK: 3
+    print(set0.n_inputs)
+    # CHECK: 1
+    print(set0.n_equalities)
+    # CHECK: 2
+    print(set0.n_inequalities)
+
+    # CHECK: 3
+    print(len(set0.constraints))
+
+    # CHECK-DAG: d0 - d1 == 0
+    # CHECK-DAG: s0 - 42 >= 0
+    # CHECK-DAG: -d0 + s0 >= 0
+    for cstr in set0.constraints:
+      print(cstr.expr, end='')
+      print(" == 0" if cstr.is_eq else " >= 0")
+
+run(testIntegerSetProperties)