[Mlir-commits] [mlir] [mlir] Add inferContractionDims util for indexing map inputs (PR #76081)

[llvmlistbot at llvm.org]

================
@@ -201,11 +193,48 @@ findPermutationsIndexingOperand(LinalgOp linalgOp, OpOperand *opOperand,
   return res;
 }
 
+/// Given a `linalgOp` and one of its `opOperand`, returns the positions of the
+/// iterators of type `iter` that index the `opOperand` as a permutation.
+/// This is useful to infer various subcomputations on a given `linalgOp`.
+/// This is performed by looking up each result in the matching indexing map and
+/// determining whether:
+///   - It is a single AffineDimExpr.
+///   - It is the only result involving this AffineDimExpr.
+static llvm::SmallDenseSet<int64_t>
+findPermutationsIndexingOperand(LinalgOp linalgOp, OpOperand *opOperand,
+                                utils::IteratorType iter) {
+  assert(linalgOp == opOperand->getOwner() && "expected linalgOp owner");
+  return findPermutationsIndexingOperandImpl(
+      linalgOp.getMatchingIndexingMap(opOperand),
+      linalgOp.getIteratorTypesArray(), iter);
+}
+
+static llvm::SmallDenseSet<int64_t>
+findPermutationsIndexingOperand(AffineMap indexingMap,
----------------
Max191 wrote:

I personally prefer not to call the impl function and just use overloads, but I can change it if you prefer me to directly call the impl given that it is just a static helper.

https://github.com/llvm/llvm-project/pull/76081