[MLIR][Python] Add bindings for PDL constraint function registering (#160520)

This is a follow-up to #159926.

That PR (#159926) exposed native rewrite function registration in PDL
through the C API and Python, enabling use with
`pdl.apply_native_rewrite`.

In this PR, we add support for native constraint functions in PDL via
`pdl.apply_native_constraint`, further completing the PDL API.

Author: PragmaTwice

mlir/include/mlir-c/Rewrite.h

@@ -375,6 +375,20 @@ MLIR_CAPI_EXPORTED void mlirPDLPatternModuleRegisterRewriteFunction(
     MlirPDLPatternModule pdlModule, MlirStringRef name,
     MlirPDLRewriteFunction rewriteFn, void *userData);
 
+/// This function type is used as callbacks for PDL native constraint functions.
+/// Input values can be accessed by `values` with its size `nValues`;
+/// output values can be added into `results` by `mlirPDLResultListPushBack*`
+/// APIs. And the return value indicates whether the constraint holds.
+typedef MlirLogicalResult (*MlirPDLConstraintFunction)(
+    MlirPatternRewriter rewriter, MlirPDLResultList results, size_t nValues,
+    MlirPDLValue *values, void *userData);
+
+/// Register a constraint function into the given PDL pattern module.
+/// `userData` will be provided as an argument to the constraint function.
+MLIR_CAPI_EXPORTED void mlirPDLPatternModuleRegisterConstraintFunction(
+    MlirPDLPatternModule pdlModule, MlirStringRef name,
+    MlirPDLConstraintFunction constraintFn, void *userData);
+
 #endif // MLIR_ENABLE_PDL_IN_PATTERNMATCH
 
 #undef DEFINE_C_API_STRUCT

mlir/lib/Bindings/Python/Rewrite.cpp

@@ -40,6 +40,15 @@ static nb::object objectFromPDLValue(MlirPDLValue value) {
   throw std::runtime_error("unsupported PDL value type");
 }
 
+static std::vector<nb::object> objectsFromPDLValues(size_t nValues,
+                                                    MlirPDLValue *values) {
+  std::vector<nb::object> args;
+  args.reserve(nValues);
+  for (size_t i = 0; i < nValues; ++i)
+    args.push_back(objectFromPDLValue(values[i]));
+  return args;
+}
+
 // Convert the Python object to a boolean.
 // If it evaluates to False, treat it as success;
 // otherwise, treat it as failure.
@@ -74,11 +83,22 @@ class PyPDLPatternModule {
            size_t nValues, MlirPDLValue *values,
            void *userData) -> MlirLogicalResult {
           nb::handle f = nb::handle(static_cast<PyObject *>(userData));
-          std::vector<nb::object> args;
-          args.reserve(nValues);
-          for (size_t i = 0; i < nValues; ++i)
-            args.push_back(objectFromPDLValue(values[i]));
-          return logicalResultFromObject(f(rewriter, results, args));
+          return logicalResultFromObject(
+              f(rewriter, results, objectsFromPDLValues(nValues, values)));
+        },
+        fn.ptr());
+  }
+
+  void registerConstraintFunction(const std::string &name,
+                                  const nb::callable &fn) {
+    mlirPDLPatternModuleRegisterConstraintFunction(
+        get(), mlirStringRefCreate(name.data(), name.size()),
+        [](MlirPatternRewriter rewriter, MlirPDLResultList results,
+           size_t nValues, MlirPDLValue *values,
+           void *userData) -> MlirLogicalResult {
+          nb::handle f = nb::handle(static_cast<PyObject *>(userData));
+          return logicalResultFromObject(
+              f(rewriter, results, objectsFromPDLValues(nValues, values)));
         },
         fn.ptr());
   }
@@ -199,6 +219,13 @@ void mlir::python::populateRewriteSubmodule(nb::module_ &m) {
              const nb::callable &fn) {
             self.registerRewriteFunction(name, fn);
           },
+          nb::keep_alive<1, 3>())
+      .def(
+          "register_constraint_function",
+          [](PyPDLPatternModule &self, const std::string &name,
+             const nb::callable &fn) {
+            self.registerConstraintFunction(name, fn);
+          },
           nb::keep_alive<1, 3>());
 #endif // MLIR_ENABLE_PDL_IN_PATTERNMATCH
   nb::class_<PyFrozenRewritePatternSet>(m, "FrozenRewritePatternSet")

mlir/lib/CAPI/Transforms/Rewrite.cpp

@@ -398,21 +398,41 @@ void mlirPDLResultListPushBackAttribute(MlirPDLResultList results,
   unwrap(results)->push_back(unwrap(value));
 }
 
+inline std::vector<MlirPDLValue> wrap(ArrayRef<PDLValue> values) {
+  std::vector<MlirPDLValue> mlirValues;
+  mlirValues.reserve(values.size());
+  for (auto &value : values) {
+    mlirValues.push_back(wrap(&value));
+  }
+  return mlirValues;
+}
+
 void mlirPDLPatternModuleRegisterRewriteFunction(
     MlirPDLPatternModule pdlModule, MlirStringRef name,
     MlirPDLRewriteFunction rewriteFn, void *userData) {
   unwrap(pdlModule)->registerRewriteFunction(
       unwrap(name),
       [userData, rewriteFn](PatternRewriter &rewriter, PDLResultList &results,
                             ArrayRef<PDLValue> values) -> LogicalResult {
-        std::vector<MlirPDLValue> mlirValues;
-        mlirValues.reserve(values.size());
-        for (auto &value : values) {
-          mlirValues.push_back(wrap(&value));
-        }
+        std::vector<MlirPDLValue> mlirValues = wrap(values);
         return unwrap(rewriteFn(wrap(&rewriter), wrap(&results),
                                 mlirValues.size(), mlirValues.data(),
                                 userData));
       });
 }
+
+void mlirPDLPatternModuleRegisterConstraintFunction(
+    MlirPDLPatternModule pdlModule, MlirStringRef name,
+    MlirPDLConstraintFunction constraintFn, void *userData) {
+  unwrap(pdlModule)->registerConstraintFunction(
+      unwrap(name),
+      [userData, constraintFn](PatternRewriter &rewriter,
+                               PDLResultList &results,
+                               ArrayRef<PDLValue> values) -> LogicalResult {
+        std::vector<MlirPDLValue> mlirValues = wrap(values);
+        return unwrap(constraintFn(wrap(&rewriter), wrap(&results),
+                                   mlirValues.size(), mlirValues.data(),
+                                   userData));
+      });
+}
 #endif // MLIR_ENABLE_PDL_IN_PATTERNMATCH