[Mlir-commits] [mlir] [mlir][linalg] Add support for scalable vectorization of linalg.mmt4d (PR #146531)

[Andrzej Warzyński]

================
@@ -324,13 +326,24 @@ struct VectorizationState {
   /// shape.
   SmallVector<bool> scalableVecDims;
 
-  /// Holds the active masks for permutations of the canonical vector iteration
-  /// space.
+  /// Holds the active masks for permutations of the canonical vector
+  /// iteration space.
   DenseMap<AffineMap, Value> activeMaskCache;
 
   /// Global vectorization guard for the incoming rewriter. It's initialized
   /// when the vectorization state is initialized.
   OpBuilder::InsertionGuard rewriterGuard;
+
+  /// Do all scalable vector sizes match the corresponding input dim sizes?
+  /// (tensor or memref)
+  ///
+  /// At the Tensor + MemRef levels, scalable sizes are modelled using
+  /// dynamic dimensions (i.e. `?`). In many cases these sizes result from
+  /// e.g. "scalable packing + tiling" and are known to always match the
+  /// scalable vector sizes. In such cases, masking can be safely skipped,
+  /// despite the presence of dynamic shapes. Use this flag with care and
+  /// only for cases where you are confident the assumption holds.
+  bool assumeScalableVecSizesMatchDimSize = false;
----------------
banach-space wrote:

> Could we make this generic, not only for scalable vectors?

Sorry, I over-indexed on "scalable vectors" .

> Assuming a dynamic dimension is multiple of a vector size, scalable or not, is useful in general.

For my needs ATM, **equality** (rather than **divisibility**) is sufficient. Do you reckon that we will require **divisibility** in the future? 

> However, IIUC, what's missing is something like ValueRangeAnalysis (not just for integers) because what you wonder is if the size if multiple of vscale or not?

We do actually have and use ValueRangeAnalysis for "scalable vectors", see e.g. in [Codegen/Transforms/Transforms.cpp](https://github.com/iree-org/iree/blob/a2004fcb0a73f8d43913bcc8011ee5c5f285ec27/compiler/src/iree/compiler/Codegen/Transforms/Transforms.cpp#L153-L156) in IREE. However, AFAIK, that analysis is costly and in the case of `linalg.mmt4d`, we just know up front that the assumption holds. Hence taking this approach rather through other means. But now I see that there's even more options in IREE that I should consider in the future, thanks for the pointers @hanhanW 🙏🏻 

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