From 276cbead8c909b5fcd72fa9ea5bd398dad30a74e Mon Sep 17 00:00:00 2001
From: Jose Lopes <jose.lopes@amd.com>
Date: Wed, 18 Sep 2024 10:50:06 +0100
Subject: [PATCH] Relaxe affine

---
 mlir/lib/Dialect/Affine/IR/AffineOps.cpp | 8 +++++---
 mlir/test/Dialect/Affine/invalid.mlir    | 2 --
 2 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
index 6baed92d208fb3..e5ceb122507e16 100644
--- a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
+++ b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
@@ -499,11 +499,13 @@ ParseResult mlir::affine::parseDimAndSymbolList(
 template <typename OpTy>
 static LogicalResult
 verifyDimAndSymbolIdentifiers(OpTy &op, Operation::operand_range operands,
-                              unsigned numDims) {
+                              unsigned numDims,
+                              bool allowNonAffineDimOperands = false) {
   unsigned opIt = 0;
   for (auto operand : operands) {
     if (opIt++ < numDims) {
-      if (!isValidDim(operand, getAffineScope(op)))
+      if (!isValidDim(operand, getAffineScope(op)) &&
+          !(allowNonAffineDimOperands && operand.getType().isIndex()))
         return op.emitOpError("operand cannot be used as a dimension id");
     } else if (!isValidSymbol(operand, getAffineScope(op))) {
       return op.emitOpError("operand cannot be used as a symbol");
@@ -2804,7 +2806,7 @@ LogicalResult AffineIfOp::verify() {
 
   // Verify that the operands are valid dimension/symbols.
   if (failed(verifyDimAndSymbolIdentifiers(*this, getOperands(),
-                                           condition.getNumDims())))
+                                           condition.getNumDims(), true)))
     return failure();
 
   return success();
diff --git a/mlir/test/Dialect/Affine/invalid.mlir b/mlir/test/Dialect/Affine/invalid.mlir
index 60f13102f55156..709d197bd4cf89 100644
--- a/mlir/test/Dialect/Affine/invalid.mlir
+++ b/mlir/test/Dialect/Affine/invalid.mlir
@@ -97,7 +97,6 @@ func.func @affine_if_invalid_dim(%arg : index) {
   affine.for %n0 = 0 to 7 {
     %dim = arith.addi %arg, %arg : index
 
-    // expected-error@+1 {{operand cannot be used as a dimension id}}
     affine.if #set0(%dim)[%n0] {}
   }
   return
@@ -109,7 +108,6 @@ func.func @affine_if_invalid_dim(%arg : index) {
 
 func.func @affine_if_invalid_sym() {
   affine.for %i0 = 0 to 7 {
-    // expected-error@+1 {{operand cannot be used as a symbol}}
     affine.if #set0(%i0)[%i0] {}
   }
   return