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
@@ -257,8 +257,8 @@
 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
--- a/mlir/lib/Parser/AffineParser.cpp
+++ b/mlir/lib/Parser/AffineParser.cpp
@@ -597,40 +597,42 @@
 }
 
 /// 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
--- /dev/null
+++ b/mlir/test/IR/affine-set.mlir
@@ -0,0 +1,49 @@
+// RUN: mlir-opt -allow-unregistered-dialect %s | FileCheck %s
+
+// Check if affine constraints with affine exprs on RHS can be parsed.
+
+// CHECK-DAG: #set{{[0-9]+}} = affine_set<(d0) : (d0 - 1 == 0)>
+#set0 = affine_set<(i) : (i == 1)>
+
+// CHECK-DAG: "testset0"() {set = #set{{[0-9]+}}} : () -> ()
+"testset0"() {set = #set0} : () -> ()
+
+// ---
+
+// CHECK-DAG: #set{{[0-9]+}} = affine_set<(d0)[s0, s1] : (d0 >= 0, -d0 + s0 >= 0, s0 - 5 == 0, -d0 + s1 + 1 >= 0)>
+#set1 = affine_set<(i)[N, M] : (i >= 0, N >= i, N == 5, M + 1 >= i)>
+
+// CHECK-DAG: "testset1"() {set = #set{{[0-9]+}}} : () -> ()
+"testset1"() {set = #set1} : () -> ()
+
+// ---
+
+// CHECK-DAG: #set{{[0-9]+}} = affine_set<(d0, d1)[s0] : (d0 >= 0, -d0 + s0 >= 0, d1 >= 0, d0 - d1 >= 0)>
+#set2 = affine_set<(i, j)[N] : (i >= 0, N >= i, j >= 0, i >= j)>
+
+// CHECK-DAG: "testset2"() {set = #set{{[0-9]+}}} : () -> ()
+"testset2"() {set = #set2} : () -> ()
+
+// ---
+
+// CHECK-DAG: #set{{[0-9]+}} = affine_set<(d0, d1)[s0, s1] : (-(d0 + d1 + s0 + s1) == 0, d0 + d1 - (s0 + s1) == 0)>
+#set3 = affine_set<(i0, i1)[N, M] : (0 == i0 + i1 + N + M, i0 + i1 == N + M)>
+
+// CHECK-DAG: "testset3"() {set = #set{{[0-9]+}}} : () -> ()
+"testset3"() {set = #set3} : () -> ()
+
+// ---
+
+// CHECK-DAG: #set{{[0-9]+}} = affine_set<(d0, d1)[s0, s1] : (-(d0 + d1 + s0 + s1) >= 0, d0 + d1 - (s0 + s1) >= 0)>
+#set4 = affine_set<(i0, i1)[N, M] : (0 >= i0 + i1 + N + M, i0 + i1 >= N + M)>
+
+// CHECK-DAG: "testset4"() {set = #set{{[0-9]+}}} : () -> ()
+"testset4"() {set = #set4} : () -> ()
+
+// ---
+
+// 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)>
+#set5 = 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: "testset5"() {set = #set{{[0-9]+}}} : () -> ()
+"testset5"() {set = #set5} : () -> ()
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
@@ -286,23 +286,7 @@
 
 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}}
-  }
-}
-
-// -----
-
-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 +294,7 @@
 
 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 +302,7 @@
 
 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 +310,7 @@
 // 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}}
   }
 }