review comments

Author: KoolJBlack

mlir/lib/Dialect/Vector/Transforms/VectorDropLeadUnitDim.cpp

@@ -400,15 +400,16 @@ mlir::vector::castAwayContractionLeadingOneDim(vector::ContractionOp contractOp,
         }
       }
 
-      // Check if the transpose effects outer unit dims only. Such transposes do
-      // not materially effect the underlying vector and can be omitted.
+      // Checks if only the outer, unit dimensions (of size 1) are permuted.
+      // Such transposes do not materially effect the underlying vector and can
+      // be omitted. EG: perm [1, 0, 2] applied to vector<1x1x8xi32>
       bool tranposeNonOuterUnitDims = false;
-      for (int64_t i = 0; i < (int64_t)perm.size(); ++i) {
-        if (perm[i] != i && i != (int64_t)perm.size() - 1) {
-          if (operands[it.index()].getType().cast<ShapedType>().getDimSize(i) !=
-              1) {
-            tranposeNonOuterUnitDims = true;
-          }
+      for (auto [index, dim] :
+           llvm::enumerate(ArrayRef<int64_t>(perm).drop_back(1))) {
+        if (dim != static_cast<int64_t>(index) &&
+            operands[it.index()].getType().cast<ShapedType>().getDimSize(
+                index) != 1) {
+          tranposeNonOuterUnitDims = true;
         }
       }
 

mlir/test/Dialect/Vector/vector-dropleadunitdim-transforms.mlir

@@ -170,17 +170,18 @@ func.func @cast_away_contraction_leading_one_dims_nonleadingunitdim_rank4_acctra
 // CHECK-DAG: #[[$MAP_1:.+]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
 // CHECK-DAG: #[[$MAP_2:.+]] = affine_map<(d0, d1, d2) -> (d1)>
 
-// CHECK-LABEL:   func.func @cast_away_contraction_leading_one_dims_vec_mat(
-// CHECK-SAME:                                 %[[VAL_0:.*]]: vector<1x1x8xi32>,
-// CHECK-SAME:                                 %[[VAL_1:.*]]: vector<1x8x8xi32>,
-// CHECK-SAME:                                 %[[VAL_2:.*]]: vector<1x8xi32>) -> vector<1x8xi32> {
-// CHECK:           %[[VAL_3:.*]] = vector.extract %[[VAL_0]][0] : vector<1x8xi32> from vector<1x1x8xi32>
-// CHECK:           %[[VAL_4:.*]] = vector.extract %[[VAL_2]][0] : vector<8xi32> from vector<1x8xi32>
-// CHECK:           %[[VAL_5:.*]] = vector.contract {indexing_maps = [#[[$MAP_0]], #[[$MAP_1]], #[[$MAP_2]]], iterator_types = ["parallel", "parallel", "reduction"], kind = #vector.kind<add>} %[[VAL_3]], %[[VAL_1]], %[[VAL_4]] : vector<1x8xi32>, vector<1x8x8xi32> into vector<8xi32>
-// CHECK:           %[[VAL_6:.*]] = vector.broadcast %[[VAL_5]] : vector<8xi32> to vector<1x8xi32>
-// CHECK:           return %[[VAL_6]] : vector<1x8xi32>
+// CHECK-LABEL:   func.func @cast_away_contraction_does_not_transpose_leading_unit_dims(
+// CHECK-SAME:                                 %[[LHS:.*]]: vector<1x1x8xi32>,
+// CHECK-SAME:                                 %[[RHS:.*]]: vector<1x8x8xi32>,
+// CHECK-SAME:                                 %[[ACC:.*]]: vector<1x8xi32>) -> vector<1x8xi32> {
+// CHECK:           %[[EXT_LHS:.*]] = vector.extract %[[LHS]][0] : vector<1x8xi32> from vector<1x1x8xi32>
+// CHECK:           %[[EXT_ACC:.*]] = vector.extract %[[ACC]][0] : vector<8xi32> from vector<1x8xi32>
+// CHECK:           %[[RES:.*]] = vector.contract {indexing_maps = [#[[$MAP_0]], #[[$MAP_1]], #[[$MAP_2]]], iterator_types = ["parallel", "parallel", "reduction"], kind = #vector.kind<add>} %[[EXT_LHS]], %[[RHS]], %[[EXT_ACC]] : vector<1x8xi32>, vector<1x8x8xi32> into vector<8xi32>
+// CHECK:           %[[BROADCAST_RES:.*]] = vector.broadcast %[[RES]] : vector<8xi32> to vector<1x8xi32>
+// CHECK:           return %[[BROADCAST_RES]] : vector<1x8xi32>
 // CHECK:         }
-func.func @cast_away_contraction_leading_one_dims_vec_mat(%lhs: vector<1x1x8xi32>,
+// CHECK-NOT        vector.transpose
+func.func @cast_away_contraction_does_not_transpose_leading_unit_dims(%lhs: vector<1x1x8xi32>,
                           %rhs: vector<1x8x8xi32>,
                           %acc: vector<1x8xi32>) -> vector<1x8xi32> {
   %result = vector.contract {indexing_maps = [affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>, affine_map<(d0, d1, d2, d3) -> (d0, d2, d3)>, affine_map<(d0, d1, d2, d3) -> (d1, d2)>], iterator_types = ["parallel", "parallel", "parallel", "reduction"], kind = #vector.kind<add>} %lhs, %rhs, %acc : vector<1x1x8xi32>, vector<1x8x8xi32> into vector<1x8xi32>