[mlir][vector] Project out anonymous bounds in ScalableValueBoundsConstraintSet (#96499)

If we don't eliminate these columns, then in some cases we fail to
compute a scalable bound. Test case reduced from a real-world example.

2 files changed, 30 insertions, 0 deletions

diff --git a/mlir/lib/Dialect/Vector/IR/ScalableValueBoundsConstraintSet.cpp b/mlir/lib/Dialect/Vector/IR/ScalableValueBoundsConstraintSet.cpp
index 9c36537..4a826f0 100644
--- a/mlir/lib/Dialect/Vector/IR/ScalableValueBoundsConstraintSet.cpp
+++ b/mlir/lib/Dialect/Vector/IR/ScalableValueBoundsConstraintSet.cpp
@@ -8,6 +8,7 @@
 
 #include "mlir/Dialect/Vector/IR/ScalableValueBoundsConstraintSet.h"
 #include "mlir/Dialect/Vector/IR/VectorOps.h"
+
 namespace mlir::vector {
 
 FailureOr<ConstantOrScalableBound::BoundSize>
@@ -74,6 +75,7 @@ ScalableValueBoundsConstraintSet::computeScalableBound(
     return p.first != scalableCstr.getVscaleValue() && !isStartingPoint;
   };
   scalableCstr.projectOut(projectOutFn);
+  scalableCstr.projectOutAnonymous(/*except=*/pos);
   // Also project out local variables (these are not tracked by the
   // ValueBoundsConstraintSet).
   for (unsigned i = 0, e = scalableCstr.cstr.getNumLocalVars(); i < e; ++i) {
diff --git a/mlir/test/Dialect/Vector/test-scalable-bounds.mlir b/mlir/test/Dialect/Vector/test-scalable-bounds.mlir
index 673e03f..6af904b 100644
--- a/mlir/test/Dialect/Vector/test-scalable-bounds.mlir
+++ b/mlir/test/Dialect/Vector/test-scalable-bounds.mlir
@@ -215,3 +215,31 @@ func.func @unsupported_negative_mod() {
   "test.some_use"(%bound) : (index) -> ()
   return
 }
+
+// -----
+
+// CHECK: #[[$SCALABLE_BOUND_MAP_5:.*]] = affine_map<()[s0] -> (s0 * 4)>
+
+// CHECK-LABEL: @extract_slice_loop
+//       CHECK:   %[[VSCALE:.*]] = vector.vscale
+//       CHECK:   %[[SCALABLE_BOUND:.*]] = affine.apply #[[$SCALABLE_BOUND_MAP_5]]()[%[[VSCALE]]]
+//       CHECK:   "test.some_use"(%[[SCALABLE_BOUND]]) : (index) -> ()
+
+func.func @extract_slice_loop(%tensor: tensor<1x1x3x?xf32>) {
+  %vscale = vector.vscale
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %c2 = arith.constant 2 : index
+  %c3 = arith.constant 3 : index
+  %c4 = arith.constant 4 : index
+  %cst = arith.constant 0.0 : f32
+  %c4_vscale = arith.muli %c4, %vscale : index
+  %slice = tensor.extract_slice %tensor[0, 0, 0, 0] [1, 1, 3, %c4_vscale] [1, 1, 1, 1] : tensor<1x1x3x?xf32> to tensor<1x3x?xf32>
+  %15 = scf.for %arg6 = %c0 to %c3 step %c1 iter_args(%arg = %slice) -> (tensor<1x3x?xf32>) {
+    %dim = tensor.dim %arg, %c2 : tensor<1x3x?xf32>
+    %bound = "test.reify_bound"(%dim) {type = "LB", vscale_min = 1, vscale_max = 16, scalable} : (index) -> index
+    "test.some_use"(%bound) : (index) -> ()
+    scf.yield %arg : tensor<1x3x?xf32>
+  }
+  return
+}