[Mlir-commits] [mlir] a829194 - [mlir][vector] Generalize castAwayContractionLeadingOneDim (#187312)

[llvmlistbot at llvm.org]

Author: Erick Ochoa Lopez
Date: 2026-04-23T14:40:35Z
New Revision: a829194012f28b6371e910d5f492e00ecb02f75f

URL: https://github.com/llvm/llvm-project/commit/a829194012f28b6371e910d5f492e00ecb02f75f
DIFF: https://github.com/llvm/llvm-project/commit/a829194012f28b6371e910d5f492e00ecb02f75f.diff

LOG: [mlir][vector] Generalize castAwayContractionLeadingOneDim (#187312)

This PR generalizes castAwayContractionLeadingOneDim to allow
accumulators with rank 1 to be matched.
With this generalization we allow the following contractions:

```
          %c = vector.contract {
               indexing_maps = [
                    affine_map<(d0, d1) -> (d0)>,
                    affine_map<(d0, d1) -> (d1, d0)>,
                    affine_map<(d0, d1) -> (d1)>],
               iterator_types = ["reduction", "parallel"],
               kind = #vector.kind<add>
               } %0, %1, %2 : vector<64xf32>, vector<1x64xf32> into vector<1xf32>
```

to be matched and transformed to 

```
map = affine_map<(d0) -> (d0)>
map1 = affine_map<(d0) -> ()>

    %2 = vector.contract {
            indexing_maps = [#map, #map, #map1],
            iterator_types = ["reduction"],
            kind = #vector.kind<add>
    } %arg0, %0, %1 : vector<64xf32>, vector<64xf32> into f32
```

I noticed the following chain of transformations in IREE which led me to
this fix:

```
          %23 = vector.shape_cast %22 : vector<64xf32> to vector<1x64xf32>
          %24 = vector.contract {indexing_maps = [affine_map<(d0, d1) -> (d0)>, affine_map<(d0, d1) -> (d1, d0)>, affine_map<(d0, d1) -> (d1)>], iterator_types = ["reduction", "parallel"], kind = #vector.kind<add>} %14, %23, %cst_0 : vector<64xf32>, vector<1x64xf32> into vector<1xf32>
```

ContractionOpToOuterProductOpLowering 

```
          %23 = vector.shape_cast %22 : vector<64xf32> to vector<1x64xf32>
          // transpose rhs
          %24 = vector.transpose %23, [1, 0] : vector<1x64xf32> to vector<64x1xf32> 
```

FoldTransposeShapeCast

```
          %21 = vector.shape_cast %20 : vector<64xf32> to vector<64x1xf32>
```

With this generalization, `castAwayContractionLeadingOneDim`'s will be
transformed to a reduction during ContractionOpLowering.

Assisted-by: Claude

Added: 
    

Modified: 
    mlir/lib/Dialect/Vector/Transforms/VectorDropLeadUnitDim.cpp
    mlir/test/Dialect/Vector/vector-dropleadunitdim-transforms.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/Vector/Transforms/VectorDropLeadUnitDim.cpp b/mlir/lib/Dialect/Vector/Transforms/VectorDropLeadUnitDim.cpp
index aa2dd89b182e1..fa95f96b88177 100644
--- a/mlir/lib/Dialect/Vector/Transforms/VectorDropLeadUnitDim.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/VectorDropLeadUnitDim.cpp
@@ -338,7 +338,7 @@ mlir::vector::castAwayContractionLeadingOneDim(vector::ContractionOp contractOp,
   VectorType oldAccType = dyn_cast<VectorType>(contractOp.getAccType());
   if (oldAccType == nullptr)
     return failure();
-  if (oldAccType.getRank() < 2)
+  if (oldAccType.getRank() < 1)
     return failure();
   if (oldAccType.getShape()[0] != 1)
     return failure();

diff  --git a/mlir/test/Dialect/Vector/vector-dropleadunitdim-transforms.mlir b/mlir/test/Dialect/Vector/vector-dropleadunitdim-transforms.mlir
index 9526d610e490e..aee77ce3da553 100644
--- a/mlir/test/Dialect/Vector/vector-dropleadunitdim-transforms.mlir
+++ b/mlir/test/Dialect/Vector/vector-dropleadunitdim-transforms.mlir
@@ -250,6 +250,25 @@ func.func @cast_away_contraction_does_not_transpose_leading_unit_dims(%lhs: vect
   return %result : vector<1x8xi32>
 }
 
+// -----
+
+// CHECK-DAG: #[[$map_dp0:.*]] = affine_map<(d0) -> (d0)>
+// CHECK-DAG: #[[$map_dp1:.*]] = affine_map<(d0) -> ()>
+
+// CHECK-LABEL: cast_away_contraction_leading_one_dims_to_dot_product
+//  CHECK-NEXT:   %[[R0:.+]] = vector.extract %{{.*}}[0] : vector<64xf32> from vector<1x64xf32>
+//  CHECK-NEXT:   %[[R1:.+]] = vector.extract %{{.*}}[0] : f32 from vector<1xf32>
+//  CHECK-NEXT:   %[[R2:.+]] = vector.contract {indexing_maps = [#[[$map_dp0]], #[[$map_dp0]], #[[$map_dp1]]],
+//  CHECK-SAME:   iterator_types = ["reduction"], kind = #vector.kind<add>}
+//  CHECK-SAME:   %{{.*}}, %[[R0]], %[[R1]] : vector<64xf32>, vector<64xf32> into f32
+//  CHECK-NEXT:   %[[R3:.+]] = vector.broadcast %[[R2]] : f32 to vector<1xf32>
+//  CHECK-NEXT:  return %[[R3]] : vector<1xf32>
+
+func.func @cast_away_contraction_leading_one_dims_to_dot_product(%arg0: vector<64xf32>, %arg1: vector<1x64xf32>, %arg2: vector<1xf32>) -> vector<1xf32> {
+  %0 = vector.contract {indexing_maps = [affine_map<(d0, d1) -> (d0)>, affine_map<(d0, d1) -> (d1, d0)>, affine_map<(d0, d1) -> (d1)>], iterator_types = ["reduction", "parallel"], kind = #vector.kind<add>} %arg0, %arg1, %arg2 : vector<64xf32>, vector<1x64xf32> into vector<1xf32>
+  return %0 : vector<1xf32>
+}
+
 // -----
 // CHECK-LABEL: func @cast_away_extract_strided_slice_leading_one_dims
 func.func @cast_away_extract_strided_slice_leading_one_dims(%arg0: vector<1x8x8xf16>) -> vector<1x1x8xf16> {