[Mlir-commits] [mlir] [mlir][vector][contract] Add comments and rename helper hooks (NFC) (PR #237708)

[Andrzej Warzyński]

https://github.com/banach-space updated https://github.com/llvm/llvm-project/pull/237708

>From 9d37d7155bbe46d1336adeb555eebe097ee006b1 Mon Sep 17 00:00:00 2001
From: Andrzej Warzynski <andrzej.warzynski at arm.com>
Date: Sat, 31 Oct 2026 19:52:02 +0000
Subject: [PATCH 1/2] [mlir][vector][contract] Add comments and rename helper
 hooks (NFC)

* Renames `getDimMap` to `getDimPairForIterTy` to make its purpose
  clearer.
* Adds a `ContractionOp` argument to the updated helpers to make it
  explicit that they only apply to that op.
* Reduces code duplication.
---
 mlir/lib/Dialect/Vector/IR/VectorOps.cpp | 34 +++++++++++++-----------
 1 file changed, 18 insertions(+), 16 deletions(-)

diff --git a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
index eedf483df2d31..03eee23ac28a4 100644
--- a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
+++ b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
@@ -1263,16 +1263,22 @@ static int64_t getResultIndex(AffineMap map, AffineExpr targetExpr) {
   return -1;
 }
 
+// For every contraction `dim` for which:
+//  * the corresponding iteratory type matches `targetIteratorTy`,
+// returns a pair: <lhsIdx, rhsIdx>, where:
+//  * {lhs|rhs}Idx is the index of `dim` in the (result of) indexing map for
+//    {lhs|rhs}.
 static std::vector<std::pair<int64_t, int64_t>>
-getDimMap(ArrayRef<AffineMap> indexingMaps, ArrayAttr iteratorTypes,
-          IteratorType targetIteratorType, MLIRContext *context) {
+getDimPairForIterTy(ContractionOp *contract, IteratorType targetIteratorTy) {
+  SmallVector<AffineMap, 4> indexingMaps(contract->getIndexingMapsArray());
+  ArrayAttr iteratorTypes = contract->getIteratorTypes();
   std::vector<std::pair<int64_t, int64_t>> dimMap;
   for (const auto &it : llvm::enumerate(iteratorTypes)) {
     auto iteratorType = llvm::cast<IteratorTypeAttr>(it.value()).getValue();
-    if (iteratorType != targetIteratorType)
+    if (iteratorType != targetIteratorTy)
       continue;
     // Search lhs/rhs map results for 'targetExpr'.
-    auto targetExpr = getAffineDimExpr(it.index(), context);
+    auto targetExpr = getAffineDimExpr(it.index(), contract->getContext());
     int64_t lhsDim = getResultIndex(indexingMaps[0], targetExpr);
     int64_t rhsDim = getResultIndex(indexingMaps[1], targetExpr);
     if (lhsDim >= 0 && rhsDim >= 0)
@@ -1281,6 +1287,14 @@ getDimMap(ArrayRef<AffineMap> indexingMaps, ArrayAttr iteratorTypes,
   return dimMap;
 }
 
+std::vector<std::pair<int64_t, int64_t>> ContractionOp::getContractingDimMap() {
+  return getDimPairForIterTy(this, IteratorType::reduction);
+}
+
+std::vector<std::pair<int64_t, int64_t>> ContractionOp::getBatchDimMap() {
+  return getDimPairForIterTy(this, IteratorType::parallel);
+}
+
 void ContractionOp::getIterationBounds(
     SmallVectorImpl<int64_t> &iterationBounds) {
   auto lhsShape = getLhsType().getShape();
@@ -1319,18 +1333,6 @@ void ContractionOp::getIterationIndexMap(
   }
 }
 
-std::vector<std::pair<int64_t, int64_t>> ContractionOp::getContractingDimMap() {
-  SmallVector<AffineMap, 4> indexingMaps(getIndexingMapsArray());
-  return getDimMap(indexingMaps, getIteratorTypes(), IteratorType::reduction,
-                   getContext());
-}
-
-std::vector<std::pair<int64_t, int64_t>> ContractionOp::getBatchDimMap() {
-  SmallVector<AffineMap, 4> indexingMaps(getIndexingMapsArray());
-  return getDimMap(indexingMaps, getIteratorTypes(), IteratorType::parallel,
-                   getContext());
-}
-
 std::optional<SmallVector<int64_t, 4>>
 ContractionOp::getIterationShapeToUnroll() {
   SmallVector<int64_t, 4> shape;

>From c4afa3025b4f204e1a7eb2d0f21b033880f3c527 Mon Sep 17 00:00:00 2001
From: Andrzej Warzynski <andrzej.warzynski at arm.com>
Date: Mon, 2 Nov 2026 08:04:26 +0000
Subject: [PATCH 2/2] Address nit

---
 mlir/lib/Dialect/Vector/IR/VectorOps.cpp | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
index 03eee23ac28a4..ce0a56af0648a 100644
--- a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
+++ b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
@@ -1264,18 +1264,18 @@ static int64_t getResultIndex(AffineMap map, AffineExpr targetExpr) {
 }
 
 // For every contraction `dim` for which:
-//  * the corresponding iteratory type matches `targetIteratorTy`,
+//  * the corresponding iteratory type matches `targetIteratorType`,
 // returns a pair: <lhsIdx, rhsIdx>, where:
 //  * {lhs|rhs}Idx is the index of `dim` in the (result of) indexing map for
 //    {lhs|rhs}.
 static std::vector<std::pair<int64_t, int64_t>>
-getDimPairForIterTy(ContractionOp *contract, IteratorType targetIteratorTy) {
+getDimPairForIterTy(ContractionOp *contract, IteratorType targetIteratorType) {
   SmallVector<AffineMap, 4> indexingMaps(contract->getIndexingMapsArray());
   ArrayAttr iteratorTypes = contract->getIteratorTypes();
   std::vector<std::pair<int64_t, int64_t>> dimMap;
   for (const auto &it : llvm::enumerate(iteratorTypes)) {
     auto iteratorType = llvm::cast<IteratorTypeAttr>(it.value()).getValue();
-    if (iteratorType != targetIteratorTy)
+    if (iteratorType != targetIteratorType)
       continue;
     // Search lhs/rhs map results for 'targetExpr'.
     auto targetExpr = getAffineDimExpr(it.index(), contract->getContext());