fixup! fixup! [mlir][linalg] Add support for scalable vectorization of linalg.mmt4d

Fix the condition that checks whether masks are needed, fix test

Author: banach-space

mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp

@@ -482,18 +482,19 @@ Value VectorizationState::getOrCreateMaskFor(
   }
 
   if (assumeDynamicDimsMatchVecSizes) {
-    // Given that all _scalable vector sizes_ match the corresponding
-    // memref/tensor dim sizes, masking can be skipped provided that:
-    // * all vector sizes corresponding to dynamic dims are scalable.
-    if (llvm::all_of(llvm::zip(permutedStaticSizes, maskType.getScalableDims()),
+    // While we can _assume_ that for dynamic dim sizes the corresponding
+    // vector sizes match, we still need to check the static dim sizes to be
+    // 100% sure that masking is indeed not required.
+    if (llvm::all_of(llvm::zip(permutedStaticSizes, maskType.getShape()),
                      [](auto it) {
                        return std::get<0>(it) == ShapedType::kDynamic
-                                  ? std::get<1>(it)
-                                  : false;
-                     }))
+                                  ? true
+                                  : std::get<0>(it) == std::get<1>(it);
+                     })) {
       LDBG("Masking is not needed for masking map: " << maskingMap << "\n");
-    activeMaskCache[maskingMap] = Value();
-    return Value();
+      activeMaskCache[maskingMap] = Value();
+      return Value();
+    }
   }
 
   // Permute the iteration space value sizes to compute the mask upper bounds.

mlir/test/Dialect/Linalg/vectorization/linalg-ops.mlir

@@ -928,7 +928,7 @@ func.func @mmt4d_scalable_with_assume(%A: memref<16x16x8x1xf32>, %B: memref<16x1
 module attributes {transform.with_named_sequence} {
   transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
     %mmt4d = transform.structured.match ops{["linalg.mmt4d"]} in %arg1 : (!transform.any_op) -> !transform.any_op
-    transform.structured.vectorize %mmt4d vector_sizes [16, 16, 16, 8, [4], 1] {assume_scalable_sizes_match_dim_size} : !transform.any_op
+    transform.structured.vectorize %mmt4d vector_sizes [16, 16, 16, 8, [4], 1] {assume_dynamic_dims_match_vec_sizes} : !transform.any_op
     transform.yield
   }
 }