[Mlir-commits] [mlir] 3c043ab - [MLIR][Affine] Allow `<=` in IntegerSet constraints

[llvmlistbot at llvm.org]

Author: Groverkss
Date: 2022-07-05T12:17:31+01:00
New Revision: 3c043ab007a1ccd59ca401074b6edd6cb67a3e77

URL: https://github.com/llvm/llvm-project/commit/3c043ab007a1ccd59ca401074b6edd6cb67a3e77
DIFF: https://github.com/llvm/llvm-project/commit/3c043ab007a1ccd59ca401074b6edd6cb67a3e77.diff

LOG: [MLIR][Affine] Allow `<=` in IntegerSet constraints

This patch extends the affine parser to allow affine constraints with `<=`.
This is useful in writing unittests for Presburger library and test in general.

The internal storage and printing of IntegerSet is still in the original format.

Reviewed By: bondhugula

Differential Revision: https://reviews.llvm.org/D129046

Added: 
    

Modified: 
    mlir/docs/Dialects/Affine.md
    mlir/lib/Parser/AffineParser.cpp
    mlir/test/IR/affine-set.mlir
    mlir/test/IR/invalid.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/docs/Dialects/Affine.md b/mlir/docs/Dialects/Affine.md
index 2b8ff57f2b924..e3ea01bec800b 100644
--- a/mlir/docs/Dialects/Affine.md
+++ b/mlir/docs/Dialects/Affine.md
@@ -258,6 +258,7 @@ Syntax of affine constraints:
 
 ```
 affine-constraint ::= affine-expr `>=` `affine-expr`
+                    | affine-expr `<=` `affine-expr`
                     | affine-expr `==` `affine-expr`
 affine-constraint-conjunction ::= affine-constraint (`,` affine-constraint)*
 ```

diff  --git a/mlir/lib/Parser/AffineParser.cpp b/mlir/lib/Parser/AffineParser.cpp
index 88800cf6e373e..7fa054bf3913a 100644
--- a/mlir/lib/Parser/AffineParser.cpp
+++ b/mlir/lib/Parser/AffineParser.cpp
@@ -598,6 +598,7 @@ ParseResult AffineParser::parseAffineMapRange(unsigned numDims,
 
 /// Parse an affine constraint.
 ///  affine-constraint ::= affine-expr `>=` `affine-expr`
+///                      | affine-expr `<=` `affine-expr`
 ///                      | affine-expr `==` `affine-expr`
 ///
 /// The constraint is normalized to
@@ -622,6 +623,15 @@ AffineExpr AffineParser::parseAffineConstraint(bool *isEq) {
     return lhsExpr - rhsExpr;
   }
 
+  // affine-constraint ::= `affine-expr` `<=` `affine-expr`
+  if (consumeIf(Token::less) && consumeIf(Token::equal)) {
+    AffineExpr rhsExpr = parseAffineExpr();
+    if (!rhsExpr)
+      return nullptr;
+    *isEq = false;
+    return rhsExpr - lhsExpr;
+  }
+
   // affine-constraint ::= `affine-expr` `==` `affine-expr`
   if (consumeIf(Token::equal) && consumeIf(Token::equal)) {
     AffineExpr rhsExpr = parseAffineExpr();

diff  --git a/mlir/test/IR/affine-set.mlir b/mlir/test/IR/affine-set.mlir
index 8532a398affe6..c719f206b6a2f 100644
--- a/mlir/test/IR/affine-set.mlir
+++ b/mlir/test/IR/affine-set.mlir
@@ -41,6 +41,23 @@
 // CHECK-DAG: #set{{[0-9]+}} = affine_set<(d0, d1, d2, d3) : ((d0 + d1) mod 2 - (d2 + d3) floordiv 2 == 0, d0 mod 2 + d1 mod 2 - (d2 + d3 + d2) >= 0)>
 #set12 = affine_set<(d0, d1, r0, r1) : ((d0 + d1) mod 2 == (r0 + r1) floordiv 2, ((d0) mod 2) + ((d1) mod 2) >= (r0 + r1) + r0)>
 
+// Check if affine constraints with <= can be parsed.
+
+// CHECK-DAG: #set{{[0-9]+}} = affine_set<(d0) : (-d0 >= 0)>
+#set13 = affine_set<(i) : (i <= 0)>
+
+// CHECK-DAG: #set{{[0-9]+}} = affine_set<(d0) : (d0 >= 0, -d0 + 10 >= 0)>
+#set14 = affine_set<(i) : (i >= 0, i <= 10)>
+
+// CHECK-DAG: #set{{[0-9]+}} = affine_set<(d0, d1)[s0, s1] : (d0 >= 0, -d0 + s0 >= 0, d1 >= 0, -d1 + s1 >= 0)>
+#set15 = affine_set<(i, j)[N, M] : (i >= 0, i <= N, j >= 0, j <= M)>
+
+// CHECK-DAG: #set{{[0-9]+}} = affine_set<(d0)[s0] : (d0 >= 0, -d0 + s0 >= 0, -d0 + s0 floordiv 2 >= 0)>
+#set16 = affine_set<(i)[N] : (i >= 0, i <= N, i <= N floordiv 2)>
+
+// CHECK-DAG: #set{{[0-9]+}} = affine_set<(d0, d1, d2, d3, d4) : (-d0 >= 0, -d1 >= 0, -d2 >= 0, -d3 >= 0, -d4 >= 0)>
+#set17 = affine_set<(i0, i1, i2, i3, i4)[] : (i0 <= 0, i1 <= 0, i2 <= 0, i3 <= 0, i4 <= 0)>
+
 // CHECK-DAG: "testset0"() {set = #set{{[0-9]+}}} : () -> ()
 "testset0"() {set = #set0} : () -> ()
 
@@ -79,3 +96,18 @@
 
 // CHECK-DAG: "testset12"() {set = #set{{[0-9]+}}} : () -> ()
 "testset12"() {set = #set12} : () -> ()
+
+// CHECK-DAG: "testset13"() {set = #set{{[0-9]+}}} : () -> ()
+"testset13"() {set = #set13} : () -> ()
+
+// CHECK-DAG: "testset14"() {set = #set{{[0-9]+}}} : () -> ()
+"testset14"() {set = #set14} : () -> ()
+
+// CHECK-DAG: "testset15"() {set = #set{{[0-9]+}}} : () -> ()
+"testset15"() {set = #set15} : () -> ()
+
+// CHECK-DAG: "testset16"() {set = #set{{[0-9]+}}} : () -> ()
+"testset16"() {set = #set16} : () -> ()
+
+// CHECK-DAG: "testset17"() {set = #set{{[0-9]+}}} : () -> ()
+"testset17"() {set = #set17} : () -> ()

diff  --git a/mlir/test/IR/invalid.mlir b/mlir/test/IR/invalid.mlir
index 2a892a8b5fcec..7d111c9b69be6 100644
--- a/mlir/test/IR/invalid.mlir
+++ b/mlir/test/IR/invalid.mlir
@@ -300,14 +300,6 @@ func.func @invalid_if_conditional3() {
 
 // -----
 
-func.func @invalid_if_conditional5() {
-  affine.for %i = 1 to 10 {
-    affine.if affine_set<(i)[N] : (i <= 0)> // expected-error {{expected '== affine-expr' or '>= affine-expr' at end of affine constraint}}
-  }
-}
-
-// -----
-
 func.func @invalid_if_conditional6() {
   affine.for %i = 1 to 10 {
     affine.if affine_set<(i) : (i)> // expected-error {{expected '== affine-expr' or '>= affine-expr' at end of affine constraint}}