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

Can we give it a try? I was mainly worried about adding new API (`PopulateCustomValueBoundsFn`) to `ValueBoundsOpInterface` when it is not needed.

Setting up the RTTI should be pretty easy, you can copy this: https://github.com/llvm/llvm-project/commit/6a43523f9cdb9a364a06bf825ac67c4478b759f0

It would indeed only make sense for `ScalableValueBoundsConstraintSet`, but I don't see a problem with that.


> I don't think `ValueBoundsConstraintSet::computeDependentBound()` does exactly what I want here. First, I'm not sure I want to stop at the first `vscale` found, also `computeDependentBound()` does not project out the `valueDim`, which in my experiments gave `affine_maps` that take two input symbols (presumably one for `vscale` and the `valueDim`), which did not usefully describe the bound.
> 
> Only once I eliminated all variables except `vscale` did I get the bounds I was interested in.

At the moment, `stopFunction` controls how far the traversal goes and also what SSA values to project out. It looks like you need some more fine-grained control here: `computeScalableBound` adds constraints for `vscale`, but it does not project `vscale` out. I don't really mind having a bit of duplicated functionality in `computeScalableBound`; as you said it's quite isolated. I can look into adding the missing configuration options in a follow-up.


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