[Mlir-commits] [mlir] [mlir][Vector] Add utility for computing scalable value bounds (PR #83876)

[Matthias Springer]

================
@@ -300,3 +301,132 @@ vector::createUnrollIterator(VectorType vType, int64_t targetRank) {
   shapeToUnroll = shapeToUnroll.slice(0, firstScalableDim);
   return StaticTileOffsetRange(shapeToUnroll, /*unrollStep=*/1);
 }
+
+FailureOr<vector::ConstantOrScalableBound::BoundSize>
+vector::ConstantOrScalableBound::getSize() const {
+  if (map.isSingleConstant())
+    return BoundSize{map.getSingleConstantResult(), /*scalable=*/false};
+  if (map.getNumResults() != 1 || map.getNumInputs() != 1)
+    return failure();
+  auto binop = dyn_cast<AffineBinaryOpExpr>(map.getResult(0));
+  if (!binop || binop.getKind() != AffineExprKind::Mul)
+    return failure();
+  auto matchConstant = [&](AffineExpr expr, int64_t &constant) -> bool {
+    if (auto cst = dyn_cast<AffineConstantExpr>(expr)) {
+      constant = cst.getValue();
+      return true;
+    }
+    return false;
+  };
+  // Match `s0 * cst` or `cst * s0`:
+  int64_t cst = 0;
+  auto lhs = binop.getLHS();
+  auto rhs = binop.getRHS();
+  if ((matchConstant(lhs, cst) && isa<AffineSymbolExpr>(rhs)) ||
+      (matchConstant(rhs, cst) && isa<AffineSymbolExpr>(lhs))) {
+    return BoundSize{cst, /*scalable=*/true};
+  }
+  return failure();
+}
+
+namespace {
+struct ScalableValueBoundsConstraintSet : public ValueBoundsConstraintSet {
----------------
matthias-springer wrote:

I'm trying to see if we can reuse more functionality from `ValueBoundsConstraintSet`, so less has to be reimplemented here.

One thing that stands out to me in this PR is that `vector.vscale` does not implement `ValueBoundsOpInterface`.

Here's an idea that could work:
- `ScalableValueBoundsConstraintSet` defines 3 fields: `vscaleMin`, `vscaleMax` and `vscale`.
- `vscaleMin` and `vscaleMax` are initialized in the constructor of `ScalableValueBoundsConstraintSet`.
- `vscale` is first seen `vscale` SSA value. "empty" value in the beginning.
- The `ValueBoundsOpInterface::populateBoundsForIndexValue` implementation dyn_casts `cstr` to `ScalableValueBoundsConstraintSet &`. If that succeeds, we can get `vscaleMin`, `vscaleMax` and `vscale` and implement the same logic that you have here without the `PopulateCustomValueBoundsFn` callback.
- At that point, we may just be able to call `ValueBoundsConstraintSet::computeDependentBound`, which gives you a bound in which only certain SSA values (`vscale`) are allowed to appear.


https://github.com/llvm/llvm-project/pull/83876