[MLIR][Affine] Allow affine-expr on RHS in IntegerSet

[Groverkss]

Summary

Currently, the parser for IntegerSet, only allows constraints like:

affine-constraint ::= affine-expr `>=` `0`
                    | affine-expr `==` `0`

This form is sometimes unreadable and painful to use when writing unittests
for Presburger library and tests in general.

This patch extends the parser to allow affine constraints with affine-expr on
the RHS:

affine-constraint ::= affine-expr `>=` `affine-expr`
                    | affine-expr `==` `affine-expr`

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

[bondhugula]

This would be a great improvement -- thanks! LGTM with some comments.

mlir/lib/Parser/AffineParser.cpp

600

You'll need to update the documentation files (*.md) as well. (See Affine.md.)

mlir/test/IR/affine-set.mlir

4

The name of this file doesn't reflect these specific test cases related to checking affine expressions on the RHS. You could just move the other existing test cases on affine integer sets all into this file.

12

Isn't it five hyphens?! You don't even have a split-input-file on top. It looks like these are all dead lines.

mlir/test/IR/invalid.mlir

297

Can you add a similar negative test case for (i == )>?

[Groverkss]

Address Uday's comments

mlir/lib/Parser/AffineParser.cpp

600

I think I already did this.

mlir/test/IR/affine-set.mlir

4

Most test cases for affine_set involve sets which are used in affine.if as well, so it is hard to move them without duplicating tests.

I have a patch, which will add more tests to this file later.

I have also added some more tests to check some corner cases for affine_set since affine_set has very less tests already.

mlir/docs/Dialects/Affine.md

 classes: the set's dimension identifiers, and the set's symbolic identifiers.
 The set is viewed as being parametric on its symbolic identifiers. In the
 syntax, the list of set's dimension identifiers are enclosed in parentheses
 while its symbols are enclosed in square brackets.
 
 Syntax of affine constraints:
 
 ```
-affine-constraint ::= affine-expr `>=` `0`
-                    | affine-expr `==` `0`
+affine-constraint ::= affine-expr `>=` `affine-expr`
+                    | affine-expr `==` `affine-expr`
 affine-constraint-conjunction ::= affine-constraint (`,` affine-constraint)*
 ```
 
 Integer sets may be defined inline at the point of use, or may be hoisted to the
 top of the file and given a name with an integer set definition, and used by
 name.
 
 ```

mlir/lib/Parser/AffineParser.cpp

if (parseCommaSeparatedList(Delimiter::Paren, parseElt,
     return failure();
 
   // Parsed a valid affine map.
   result = AffineMap::get(numDims, numSymbols, exprs, getContext());
   return success();
 }
 
 /// Parse an affine constraint.
+///  affine-constraint ::= affine-expr `>=` `affine-expr`
+///                      | affine-expr `==` `affine-expr`
+///
+/// The constraint is normalized to
 ///  affine-constraint ::= affine-expr `>=` `0`
 ///                      | affine-expr `==` `0`
+/// before returning.
 ///
 /// isEq is set to true if the parsed constraint is an equality, false if it
 /// is an inequality (greater than or equal).
 ///
 AffineExpr AffineParser::parseAffineConstraint(bool *isEq) {
-  AffineExpr expr = parseAffineExpr();
-  if (!expr)
+  AffineExpr lhsExpr = parseAffineExpr();
+  if (!lhsExpr)
     return nullptr;
 
-  if (consumeIf(Token::greater) && consumeIf(Token::equal) &&
-      getToken().is(Token::integer)) {
-    auto dim = getToken().getUnsignedIntegerValue();
-    if (dim && *dim == 0) {
-      consumeToken(Token::integer);
-      *isEq = false;
-      return expr;
-    }
-    return emitError("expected '0' after '>='"), nullptr;
+  // affine-constraint ::= `affine-expr` `>=` `affine-expr`
+  if (consumeIf(Token::greater) && consumeIf(Token::equal)) {
+    AffineExpr rhsExpr = parseAffineExpr();
+    if (!rhsExpr)
+      return nullptr;
+    return lhsExpr - rhsExpr;
   }
 
-  if (consumeIf(Token::equal) && consumeIf(Token::equal) &&
-      getToken().is(Token::integer)) {
-    auto dim = getToken().getUnsignedIntegerValue();
-    if (dim && *dim == 0) {
-      consumeToken(Token::integer);
-      *isEq = true;
-      return expr;
-    }
-    return emitError("expected '0' after '=='"), nullptr;
+  // affine-constraint ::= `affine-expr` `==` `affine-expr`
+  if (consumeIf(Token::equal) && consumeIf(Token::equal)) {
+    AffineExpr rhsExpr = parseAffineExpr();
+    if (!rhsExpr)
+      return nullptr;
+    return lhsExpr - rhsExpr;
   }
 
-  return emitError("expected '== 0' or '>= 0' at end of affine constraint"),
+  return emitError("expected '== affine-expr' or '>= affine-expr' at end of "
+                   "affine constraint"),
          nullptr;
 }
 
 /// Parse the constraints that are part of an integer set definition.
 ///  integer-set-inline
 ///                ::= dim-and-symbol-id-lists `:`
 ///                '(' affine-constraint-conjunction? ')'
 ///  affine-constraint-conjunction ::= affine-constraint (`,`

mlir/test/IR/affine-set.mlir

+// RUN: mlir-opt -allow-unregistered-dialect %s | FileCheck %s
+
+// CHECK-DAG: #set{{[0-9]+}} = affine_set<() : (0 == 0)>
+#set0 = affine_set<() : ()>
+
+// CHECK-DAG: #set{{[0-9]+}} = affine_set<(d0) : (d0 == 0)>
+#set1 = affine_set<(i) : (i == 0)>
+
+// CHECK-DAG: #set{{[0-9]+}} = affine_set<(d0) : (d0 >= 0)>
+#set2 = affine_set<(i) : (i >= 0)>
+
+// CHECK-DAG: #set{{[0-9]+}} = affine_set<(d0, d1) : (d0 >= 0, d1 >= 0)>
+#set3 = affine_set<(i, j) : (i >= 0, j >= 0)>
+
+// CHECK-DAG: #set{{[0-9]+}} = affine_set<(d0, d1) : (d0 == 0, d1 >= 0)>
+#set4 = affine_set<(i, j) : (i == 0, j >= 0)>
+
+// CHECK-DAG: #set{{[0-9]+}} = affine_set<(d0, d1)[s0, s1] : (d0 * 2 + s0 >= 0, d1 * 2 + s1 >= 0, d0 >= 0, d1 >= 0, s0 == 0, s1 == 0)>
+#set5 = affine_set<(i, j)[N, M] : (i * 2 + N >= 0, j * 2 + M >= 0, i >= 0, j >= 0, N == 0, M == 0)>
+
+// CHECK-DAG: #set{{[0-9]+}} = affine_set<(d0, d1)[s0, s1] : (-d0 + s0 >= 0, -d1 + s1 >= 0, d0 >= 0, d1 >= 0)>
+#set6 = affine_set<(i, j)[N, M] : (N - i >= 0, M - j >= 0, i >= 0, j >= 0)>
+
+// Check if affine constraints with affine exprs on RHS can be parsed.
+
+// CHECK-DAG: #set{{[0-9]+}} = affine_set<(d0) : (d0 - 1 == 0)>
+#set7 = affine_set<(i) : (i == 1)>
+
+// CHECK-DAG: #set{{[0-9]+}} = affine_set<(d0)[s0, s1] : (d0 >= 0, -d0 + s0 >= 0, s0 - 5 == 0, -d0 + s1 + 1 >= 0)>
+#set8 = affine_set<(i)[N, M] : (i >= 0, N >= i, N == 5, M + 1 >= i)>
+
+// CHECK-DAG: #set{{[0-9]+}} = affine_set<(d0, d1)[s0] : (d0 >= 0, -d0 + s0 >= 0, d1 >= 0, d0 - d1 >= 0)>
+#set9 = affine_set<(i, j)[N] : (i >= 0, N >= i, j >= 0, i >= j)>
+
+// CHECK-DAG: #set{{[0-9]+}} = affine_set<(d0, d1)[s0, s1] : (-(d0 + d1 + s0 + s1) == 0, d0 + d1 - (s0 + s1) == 0)>
+#set10 = affine_set<(i0, i1)[N, M] : (0 == i0 + i1 + N + M, i0 + i1 == N + M)>
+
+// CHECK-DAG: #set{{[0-9]+}} = affine_set<(d0, d1)[s0, s1] : (-(d0 + d1 + s0 + s1) >= 0, d0 + d1 - (s0 + s1) >= 0)>
+#set11 = affine_set<(i0, i1)[N, M] : (0 >= i0 + i1 + N + M, i0 + i1 >= N + M)>
+
+// 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-DAG: "testset0"() {set = #set{{[0-9]+}}} : () -> ()
+"testset0"() {set = #set0} : () -> ()
+
+// CHECK-DAG: "testset1"() {set = #set{{[0-9]+}}} : () -> ()
+"testset1"() {set = #set1} : () -> ()
+
+// CHECK-DAG: "testset2"() {set = #set{{[0-9]+}}} : () -> ()
+"testset2"() {set = #set2} : () -> ()
+
+// CHECK-DAG: "testset3"() {set = #set{{[0-9]+}}} : () -> ()
+"testset3"() {set = #set3} : () -> ()
+
+// CHECK-DAG: "testset4"() {set = #set{{[0-9]+}}} : () -> ()
+"testset4"() {set = #set4} : () -> ()
+
+// CHECK-DAG: "testset5"() {set = #set{{[0-9]+}}} : () -> ()
+"testset5"() {set = #set5} : () -> ()
+
+// CHECK-DAG: "testset6"() {set = #set{{[0-9]+}}} : () -> ()
+"testset6"() {set = #set6} : () -> ()
+
+// CHECK-DAG: "testset7"() {set = #set{{[0-9]+}}} : () -> ()
+"testset7"() {set = #set7} : () -> ()
+
+// CHECK-DAG: "testset8"() {set = #set{{[0-9]+}}} : () -> ()
+"testset8"() {set = #set8} : () -> ()
+
+// CHECK-DAG: "testset9"() {set = #set{{[0-9]+}}} : () -> ()
+"testset9"() {set = #set9} : () -> ()
+
+// CHECK-DAG: "testset10"() {set = #set{{[0-9]+}}} : () -> ()
+"testset10"() {set = #set10} : () -> ()
+
+// CHECK-DAG: "testset11"() {set = #set{{[0-9]+}}} : () -> ()
+"testset11"() {set = #set11} : () -> ()
+
+// CHECK-DAG: "testset12"() {set = #set{{[0-9]+}}} : () -> ()
+"testset12"() {set = #set12} : () -> ()

mlir/test/IR/invalid.mlir

func.func @non_operation() {
   // expected-error@+1 {{custom op 'asd' is unknown (tried 'func.asd' as well)}}
   asd
 }
 
 // -----
 
 func.func @invalid_if_conditional2() {
   affine.for %i = 1 to 10 {
-    affine.if affine_set<(i)[N] : (i >= )>  // expected-error {{expected '== 0' or '>= 0' at end of affine constraint}}
+    affine.if affine_set<(i)[N] : (i >= )>  // expected-error {{expected affine expression}}
   }
 }
 
 // -----
 
 func.func @invalid_if_conditional3() {
   affine.for %i = 1 to 10 {
-    affine.if affine_set<(i)[N] : (i == 1)> // expected-error {{expected '0' after '=='}}
-  }
-}
-
-// -----
-
-func.func @invalid_if_conditional4() {
-  affine.for %i = 1 to 10 {
-    affine.if affine_set<(i)[N] : (i >= 2)> // expected-error {{expected '0' after '>='}}
+    affine.if affine_set<(i)[N] : (i == )>  // expected-error {{expected affine expression}}
   }
 }
 
 // -----
 
 func.func @invalid_if_conditional5() {
   affine.for %i = 1 to 10 {
-    affine.if affine_set<(i)[N] : (i <= 0)> // expected-error {{expected '== 0' or '>= 0' at end of affine constraint}}
+    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 '== 0' or '>= 0' at end of affine constraint}}
+    affine.if affine_set<(i) : (i)> // expected-error {{expected '== affine-expr' or '>= affine-expr' at end of affine constraint}}
   }
 }
 
 // -----
 // TODO: support affine.if (1)?
 func.func @invalid_if_conditional7() {
   affine.for %i = 1 to 10 {
-    affine.if affine_set<(i) : (1)> // expected-error {{expected '== 0' or '>= 0' at end of affine constraint}}
+    affine.if affine_set<(i) : (1)> // expected-error {{expected '== affine-expr' or '>= affine-expr' at end of affine constraint}}
   }
 }
 
 // -----
 
 #map = affine_map<(d0) -> (%  // expected-error {{invalid SSA name}}
 
 // -----