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

[Benjamin Maxwell]

================
@@ -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 {
----------------
MacDue wrote:

I looked into the above, and I don't really think it helps much with reuse, and is a little complex to setup. To be able to `dyn_cast` the `ValueBoundConstraintSet` we'd need to implement LLVM RTTI for it, and then it still only makes sense to use the helper on `ScalableValueBoundsConstraintSet` to compute the bound. The other variants construct a non-scalable `ValueBoundsConstraintSet` internally, so would not be useful for computing bounds dependent on `vscale`. 

And like I mentioned before none of the existing bounds functions do exactly what I want here, so I think I'd still need `computeScalableBound()` (and it's accompanying helpers). 

I think it could be done, but I think (right now) all it'd really eliminate is the `PopulateCustomValueBoundsFn` callback.  

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