diff --git a/mlir/docs/Dialects/Affine.md b/mlir/docs/Dialects/Affine.md --- a/mlir/docs/Dialects/Affine.md +++ b/mlir/docs/Dialects/Affine.md @@ -258,6 +258,7 @@ ``` 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 --- a/mlir/lib/Parser/AffineParser.cpp +++ b/mlir/lib/Parser/AffineParser.cpp @@ -598,6 +598,7 @@ /// 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 @@ 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 --- 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 --- a/mlir/test/IR/invalid.mlir +++ b/mlir/test/IR/invalid.mlir @@ -300,14 +300,6 @@ // ----- -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}}