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

[Erick Ochoa Lopez]

https://github.com/amd-eochoalo created https://github.com/llvm/llvm-project/pull/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 of `castAwayContractionLeadingOneDim` gets converted into a vector.reduction

>From 3dc9bafd93b914a380b31482a58318c8215f18f4 Mon Sep 17 00:00:00 2001
From: Erick Ochoa <erick.ochoalopez at amd.com>
Date: Wed, 18 Mar 2026 11:16:45 -0400
Subject: [PATCH] [mlir][vector] Generalize castAwayContractionLeadingOneDim

---
 .../Transforms/VectorDropLeadUnitDim.cpp      |  2 +-
 .../vector-dropleadunitdim-transforms.mlir    | 19 +++++++++++++++++++
 2 files changed, 20 insertions(+), 1 deletion(-)

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> {