[Mlir-commits] [mlir] d71a8bb - [MLIR][Affine] Allow affine-expr on RHS in IntegerSet
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Sun Jul 3 08:27:02 PDT 2022
Author: Groverkss
Date: 2022-07-03T16:22:39+01:00
New Revision: d71a8bb157ea9160f090c79aaffab87080bd210b
URL: https://github.com/llvm/llvm-project/commit/d71a8bb157ea9160f090c79aaffab87080bd210b
DIFF: https://github.com/llvm/llvm-project/commit/d71a8bb157ea9160f090c79aaffab87080bd210b.diff
LOG: [MLIR][Affine] Allow affine-expr on RHS in IntegerSet
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.
Reviewed By: bondhugula
Differential Revision: https://reviews.llvm.org/D128915
Added:
mlir/test/IR/affine-set.mlir
Modified:
mlir/docs/Dialects/Affine.md
mlir/lib/Parser/AffineParser.cpp
mlir/test/IR/invalid.mlir
Removed:
################################################################################
diff --git a/mlir/docs/Dialects/Affine.md b/mlir/docs/Dialects/Affine.md
index 89d9d89964f94..2b8ff57f2b924 100644
--- a/mlir/docs/Dialects/Affine.md
+++ b/mlir/docs/Dialects/Affine.md
@@ -257,8 +257,8 @@ 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)*
```
diff --git a/mlir/lib/Parser/AffineParser.cpp b/mlir/lib/Parser/AffineParser.cpp
index 0b1797748e0d9..88800cf6e373e 100644
--- a/mlir/lib/Parser/AffineParser.cpp
+++ b/mlir/lib/Parser/AffineParser.cpp
@@ -597,40 +597,42 @@ ParseResult AffineParser::parseAffineMapRange(unsigned numDims,
}
/// 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;
+ *isEq = false;
+ 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;
+ *isEq = true;
+ 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;
}
diff --git a/mlir/test/IR/affine-set.mlir b/mlir/test/IR/affine-set.mlir
new file mode 100644
index 0000000000000..8532a398affe6
--- /dev/null
+++ b/mlir/test/IR/affine-set.mlir
@@ -0,0 +1,81 @@
+// 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} : () -> ()
diff --git a/mlir/test/IR/invalid.mlir b/mlir/test/IR/invalid.mlir
index 3a8b7911638fe..2a892a8b5fcec 100644
--- a/mlir/test/IR/invalid.mlir
+++ b/mlir/test/IR/invalid.mlir
@@ -286,7 +286,7 @@ func.func @non_operation() {
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}}
}
}
@@ -294,15 +294,7 @@ func.func @invalid_if_conditional2() {
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}}
}
}
@@ -310,7 +302,7 @@ func.func @invalid_if_conditional4() {
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}}
}
}
@@ -318,7 +310,7 @@ func.func @invalid_if_conditional5() {
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}}
}
}
@@ -326,7 +318,7 @@ func.func @invalid_if_conditional6() {
// 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}}
}
}
More information about the Mlir-commits
mailing list