fixup! [mlir][vector] Extend mask calculation for vector.contract

Author: banach-space

mlir/lib/Dialect/Vector/IR/VectorOps.cpp

@@ -912,18 +912,18 @@ Type ContractionOp::getExpectedMaskType() {
 
   unsigned numVecDims = lhsIdxMap.getNumDims();
   SmallVector<int64_t> maskShape(numVecDims, ShapedType::kDynamic);
-  SmallVector<bool> maskShapeScalabledims(numVecDims, false);
+  SmallVector<bool> maskShapeScalableDims(numVecDims, false);
 
   // Using the information in the indexing maps, extract the size of each
   // dimension in the vector.contract operation from the two input operands.
   for (auto [dimIdx, dimSize] : llvm::enumerate(lhsType.getShape())) {
     maskShape[lhsIdxMap.getDimPosition(dimIdx)] = dimSize;
-    maskShapeScalabledims[lhsIdxMap.getDimPosition(dimIdx)] =
+    maskShapeScalableDims[lhsIdxMap.getDimPosition(dimIdx)] =
         lhsType.getScalableDims()[dimIdx];
   }
   for (auto [dimIdx, dimSize] : llvm::enumerate(rhsType.getShape())) {
     maskShape[rhsIdxMap.getDimPosition(dimIdx)] = dimSize;
-    maskShapeScalabledims[rhsIdxMap.getDimPosition(dimIdx)] =
+    maskShapeScalableDims[rhsIdxMap.getDimPosition(dimIdx)] =
         rhsType.getScalableDims()[dimIdx];
   }
 
@@ -932,7 +932,7 @@ Type ContractionOp::getExpectedMaskType() {
 
   return VectorType::get(maskShape,
                          IntegerType::get(lhsType.getContext(), /*width=*/1),
-                         maskShapeScalabledims);
+                         maskShapeScalableDims);
 }
 
 SmallVector<StringRef> ContractionOp::getTraitAttrNames() {

mlir/test/Dialect/Vector/ops.mlir

@@ -989,12 +989,17 @@ func.func @vector_scalable_extract(%sv: vector<[8]xi32>) {
   indexing_maps = #matmat_accesses,
   iterator_types = ["parallel", "parallel", "reduction"]
 }
-func.func @matmul_masked_scalable(%arg0: vector<3x4xf32>,
-                                    %arg1: vector<4x[8]xf32>,
-                                    %arg2: vector<3x[8]xf32>,
-                                    %m : vector<3x[8]x4xi1>) -> vector<3x[8]xf32> {
-  %0 = vector.mask %m { vector.contract #matmat_trait %arg0, %arg1, %arg2
-  : vector<3x4xf32>, vector<4x[8]xf32> into vector<3x[8]xf32> } : vector<3x[8]x4xi1> -> vector<3x[8]xf32>
+// CHECK-LABEL:   func.func @contraction_masked_scalable(
+// CHECK-SAME:    %[[A:.*]]: vector<3x4xf32>,
+// CHECK-SAME:    %[[B:.*]]: vector<4x[8]xf32>,
+// CHECK-SAME:    %[[C:.*]]: vector<3x[8]xf32>,
+// CHECK-SAME:    %[[M:.*]]: vector<3x[8]x4xi1>) -> vector<3x[8]xf32> {
+func.func @contraction_masked_scalable(%A: vector<3x4xf32>,
+                                    %B: vector<4x[8]xf32>,
+                                    %C: vector<3x[8]xf32>,
+                                    %M : vector<3x[8]x4xi1>) -> vector<3x[8]xf32> {
+ // CHECK:  vector.mask %[[M]] { vector.contract {indexing_maps = [#{{.*}}, #{{.*}}, #{{.*}}], iterator_types = ["parallel", "parallel", "reduction"], kind = #vector.kind<add>} %[[A]], %[[B]], %[[C]] : vector<3x4xf32>, vector<4x[8]xf32> into vector<3x[8]xf32> } : vector<3x[8]x4xi1> -> vector<3x[8]xf32>
+  %0 = vector.mask %M { vector.contract #matmat_trait %A, %B, %C : vector<3x4xf32>, vector<4x[8]xf32> into vector<3x[8]xf32> } 
+    : vector<3x[8]x4xi1> -> vector<3x[8]xf32>
   return %0 : vector<3x[8]xf32>
 }
-

mlir/test/Dialect/Vector/vector-contract-to-parallel-arith-transforms.mlir

@@ -60,4 +60,3 @@ transform.sequence failures(propagate) {
     transform.apply_patterns.vector.lower_contraction lowering_strategy = "parallelarith"
   } : !transform.any_op
 }
-