[Mlir-commits] [mlir] [mlir][vector] Project out anonymous bounds in ScalableValueBoundsConstraintSet (PR #96499)

[Benjamin Maxwell]

https://github.com/MacDue created https://github.com/llvm/llvm-project/pull/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.

>From 646935107be7799ced7b4e4fbb583071c7cc50d3 Mon Sep 17 00:00:00 2001
From: Benjamin Maxwell <benjamin.maxwell at arm.com>
Date: Mon, 24 Jun 2024 14:42:18 +0000
Subject: [PATCH] [mlir][vector] Project out anonymous bounds in
 ScalableValueBoundsConstraintSet

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.
---
 .../IR/ScalableValueBoundsConstraintSet.cpp   |  2 ++
 .../Dialect/Vector/test-scalable-bounds.mlir  | 28 +++++++++++++++++++
 2 files changed, 30 insertions(+)

diff --git a/mlir/lib/Dialect/Vector/IR/ScalableValueBoundsConstraintSet.cpp b/mlir/lib/Dialect/Vector/IR/ScalableValueBoundsConstraintSet.cpp
index 9c365376c84c9..4a826f04e1f1d 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 673e03f05c1b8..6af904beb660b 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
+}