[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 ↗

(On Diff #493156)

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/SCF/for-loop-canonicalization.mlir

func.func @regression_multiplication_with_sym(%A : memref<i64>) {
   scf.for %i = %c0 to %ub step %c1 {
     // CHECK: affine.min #[[$map]](%[[iv]])[%[[dummy]]]
     %1 = affine.min affine_map<(d0)[s0] -> (-(d0 * (5 ceildiv s0)) + 5, 3)>(%i)[%ub]
     %2 = arith.index_cast %1: index to i64
     memref.store %2, %A[]: memref<i64>
   }
   return
 }
+
+// -----
+
+// 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
+  }
+}