Merge pull request #353 from Xilinx/jose-relax-affine-if

Relax affine
Xilinx · Sep 19, 2024 · 9054950 · 9054950
2 parents b46c5b7 + c391a49
commit 9054950
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 5 deletions.
diff --git 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,8 @@ LogicalResult AffineIfOp::verify() {
 
   // Verify that the operands are valid dimension/symbols.
   if (failed(verifyDimAndSymbolIdentifiers(*this, getOperands(),
-                                           condition.getNumDims())))
+                                           condition.getNumDims(),
+                                           /*allowNonAffineDimOperands=*/true)))
     return failure();
 
   return success();

diff --git a/mlir/test/Dialect/Affine/invalid.mlir b/mlir/test/Dialect/Affine/invalid.mlir
@@ -96,8 +96,8 @@ func.func @affine_for_upper_bound_invalid_sym() {
 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}}
+    // Non-affine operand %dim has been made legal as input to affine.if.
+    // expected-error@+1 {{operand cannot be used as a symbol}}
     affine.if #set0(%dim)[%n0] {}
   }
   return