From 3c940451cf20e9bfbbe0a880464e960f77810be8 Mon Sep 17 00:00:00 2001
From: "Rickert, Jonas" <Jonas.Rickert@amd.com>
Date: Tue, 18 Feb 2025 02:15:44 +0000
Subject: [PATCH] When writing PDLL patterns it is often assumed that some
 basic checks are executed before constraints are called, but this is not
 always the case, as operations can be reordered in PDLInterp if there is no
 dependency between them.
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

For example:
Pdll pattern:
```
 let someOp = op<someDialect.SomeOp>(input: Value<inputTy: Type) {axis = inputAxis: Attr } -> (resTypes: TypeRange);
 let someResult = someConstraint(inputAxis);
```

If SomeOp requires axis to have a valid value, it is easy to (wrongly) assume that someConstraint always gets called with a not-null inputAxis. This is not correct.

The linearized PDLInterp (pseudo-code) could be the following:
```
%0 = pdl_interp.get_attribute "axis" of %arg0
%1 = pdl_interp.apply_constraint “someConstraint”(%0)
pdl_interp.is_not_null(%0)
pdl_interp.check_operation_name of %arg0 is "someDialect.SomeOp"
```

Note that here someConstraint can be called with a null attribute.

This commit changes the prioritization of predicates, so that constraints are run after other predicates.
```
%0 = pdl_interp.get_attribute "axis" of %arg0
pdl_interp.is_not_null(%0)
pdl_interp.check_operation_name of %arg0 is "someDialect.SomeOp"
%1 = pdl_interp.apply_constraint “someConstraint”(%0)
```
This ensures that null or operation name checks are run before constraints. This is closer to the mental model when writing PDLL patterns and should make it less likely to run into bugs caused by assuming implicit checks for not null.
---
 .../PDLToPDLInterp/PredicateTree.cpp          | 16 ++++++++++----
 .../pdl-to-pdl-interp-matcher.mlir            | 21 +++++++++++++++++++
 2 files changed, 33 insertions(+), 4 deletions(-)

diff --git a/mlir/lib/Conversion/PDLToPDLInterp/PredicateTree.cpp b/mlir/lib/Conversion/PDLToPDLInterp/PredicateTree.cpp
index 99bf73de59e37..5d249a37d50fd 100644
--- a/mlir/lib/Conversion/PDLToPDLInterp/PredicateTree.cpp
+++ b/mlir/lib/Conversion/PDLToPDLInterp/PredicateTree.cpp
@@ -768,17 +768,25 @@ struct OrderedPredicate {
   /// model.
   bool operator<(const OrderedPredicate &rhs) const {
     // Sort by:
+    // * not being a constraint. Rational: When writing constraints, it is
+    //   sometimes assumed that checks for null or operation names are executed
+    //   before the constraint. As there is no dependency between this
+    //   operation, this is not always guaranteed, which can lead to bugs if the
+    //   constraints is not checking inputs for null itself. By ordering
+    //   constraints to the end, it is assured that implicit checks are nun
+    //   before them
     // * higher first and secondary order sums
     // * lower depth
     // * lower position dependency
     // * lower predicate dependency
     // * lower tie breaking ID
     auto *rhsPos = rhs.position;
-    return std::make_tuple(primary, secondary, rhsPos->getOperationDepth(),
+    return std::make_tuple(!isa<ConstraintQuestion>(question), primary,
+                           secondary, rhsPos->getOperationDepth(),
                            rhsPos->getKind(), rhs.question->getKind(), rhs.id) >
-           std::make_tuple(rhs.primary, rhs.secondary,
-                           position->getOperationDepth(), position->getKind(),
-                           question->getKind(), id);
+           std::make_tuple(!isa<ConstraintQuestion>(rhs.question), rhs.primary,
+                           rhs.secondary, position->getOperationDepth(),
+                           position->getKind(), question->getKind(), id);
   }
 };
 
diff --git a/mlir/test/Conversion/PDLToPDLInterp/pdl-to-pdl-interp-matcher.mlir b/mlir/test/Conversion/PDLToPDLInterp/pdl-to-pdl-interp-matcher.mlir
index aeb4c7233d1ff..42fadbc3a6a90 100644
--- a/mlir/test/Conversion/PDLToPDLInterp/pdl-to-pdl-interp-matcher.mlir
+++ b/mlir/test/Conversion/PDLToPDLInterp/pdl-to-pdl-interp-matcher.mlir
@@ -488,6 +488,27 @@ module @predicate_ordering {
   }
 }
 
+// -----
+
+// CHECK-LABEL: module @predicate_ordering_attr
+module @predicate_ordering_attr {
+  // Check that the result is checked for null first, before applying the
+  // constraint. 
+
+  // CHECK: func @matcher(%[[ROOT:.*]]: !pdl.operation)
+  // CHECK:   %[[RESULT:.*]] = pdl_interp.get_attribute "attr" of %[[ROOT]]
+  // CHECK-NEXT: pdl_interp.is_not_null %[[RESULT]]
+  // CHECK: pdl_interp.apply_constraint "constraint"
+
+
+  pdl.pattern : benefit(1) {
+    %attr = attribute
+    pdl.apply_native_constraint "constraint"(%attr: !pdl.attribute)
+    pdl.apply_native_constraint "constraint"(%attr: !pdl.attribute)
+    %root = operation  "foo.op" {"attr" = %attr}
+    rewrite %root with "rewriter"
+  }
+}
 
 // -----