[mlir] use getNumDimAndSymbolVars when iterate dims/symbols of FlatAffineValueConstraints

[python3kgae]

Summary

Fixes #59443 https://github.com/llvm/llvm-project/issues/59443

getNumVars will add locals and cause out of bound access.

[springerm]

mlir/test/Dialect/Affine/transform-op-simplify-all-const-min.mlir

1

This test case should be in Dialect/SCF/for-loop-canonicalization.mlir. It has no scf.for loop, but that's the file that contains tests that exercise this code path.

[python3kgae]

Move test to correct place.

mlir/lib/Dialect/Affine/Analysis/Utils.cpp

mlir::simplifyConstrainedMinMaxOp(Operation *op,
   }
 
   // Lower and upper bound of `op` are equal. Replace `minOp` with its bound.
   AffineMap newMap = alignedBoundMap;
   SmallVector<Value> newOperands;
   unpackOptionalValues(constraints.getMaybeValues(), newOperands);
   // If dims/symbols have known constant values, use those in order to simplify
   // the affine map further.
-  for (int64_t i = 0, e = constraints.getNumVars(); i < e; ++i) {
+  for (int64_t i = 0, e = constraints.getNumDimAndSymbolVars(); i < e; ++i) {
     // Skip unused operands and operands that are already constants.
     if (!newOperands[i] || getConstantIntValue(newOperands[i]))
       continue;
     if (auto bound = constraints.getConstantBound64(IntegerPolyhedron::EQ, i)) {
       AffineExpr expr =
           i < newMap.getNumDims()
               ? builder.getAffineDimExpr(i)
               : builder.getAffineSymbolExpr(i - newMap.getNumDims());
       newMap = newMap.replace(expr, builder.getAffineConstantExpr(*bound),
                               newMap.getNumDims(), newMap.getNumSymbols());
     }
   }
   mlir::canonicalizeMapAndOperands(&newMap, &newOperands);
   return AffineValueMap(newMap, newOperands);
 }

mlir/test/Dialect/Affine/transform-op-simplify-all-const-min.mlir

+// RUN: mlir-opt %s --scf-for-loop-canonicalization | FileCheck %s
+
+// Make sure min is transformed into zero.
+// CHECK: %[[ZERO:.+]] = arith.constant 0 : index
+// CHECK: scf.index_switch %[[ZERO]] -> i1
+
+#map6 = affine_map<(d0, d1, d2) -> (d0 floordiv 64)>
+#map29 = affine_map<(d0, d1, d2) -> (d2 * 64 - 2, 5, (d1 mod 4) floordiv 8)>  
+module {
+  func.func @func1() { 
+    %true = arith.constant true 
+    %c0 = arith.constant 0 : index 
+    %c5 = arith.constant 5 : index
+    %c11 = arith.constant 11 : index 
+    %c14 = arith.constant 14 : index
+    %c15 = arith.constant 15 : index 
+    %alloc_249 = memref.alloc() : memref<7xf32>
+    %135 = affine.apply #map6(%c15, %c0, %c14) 
+    %163 = affine.min #map29(%c5, %135, %c11)
+    %196 = scf.index_switch %163 -> i1 
+    default { 
+      memref.assume_alignment %alloc_249, 1 : memref<7xf32>
+      scf.yield %true : i1
+    } 
+    return
+  } 
+}