[Mlir-commits] [mlir] e2f1d5c - [mlir][Vector] Add mapIterationSpaceDimToAllOperandDims to Linalg interface

[Diego Caballero]

Author: Diego Caballero
Date: 2023-03-29T19:53:29Z
New Revision: e2f1d5ce9147118c94e2d65ea5c53e7474e8ff4a

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

LOG: [mlir][Vector] Add mapIterationSpaceDimToAllOperandDims to Linalg interface

This is a variant of the existing `mapIterationSpaceDimToOperandDim`.
We have a local use downstream.

Reviewed By: nicolasvasilache

Differential Revision: https://reviews.llvm.org/D146857

Added: 
    

Modified: 
    mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.td

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.td b/mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.td
index 9b51228382933..ace2427f577cb 100644
--- a/mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.td
+++ b/mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.td
@@ -621,6 +621,31 @@ def LinalgStructuredInterface : OpInterface<"LinalgOp"> {
         return failure();
       }]
     >,
+    InterfaceMethod<
+      /*desc=*/[{
+        Given a dimension of the iteration space of a Linalg operation, finds
+        all the operands in the operation that are defined on such dimension.
+        Returns all the operand values found and their dimension positions in
+        `operandDimPairs`.
+      }],
+      /*retTy=*/"void",
+      /*methodName=*/"mapIterationSpaceDimToAllOperandDims",
+      /*args=*/(ins "unsigned":$dimPos,
+                    "mlir::SmallVectorImpl<std::pair<Value, unsigned>>&":$operandDimPairs),
+      /*methodBody=*/"",
+      /*defaultImplementation=*/[{
+        for (auto [i, idxMap] : llvm::enumerate($_op.getIndexingMapsArray())) {
+          if (idxMap.isProjectedPermutation()) {
+            if (auto mayOperandDim = idxMap.getResultPosition(
+                getAffineDimExpr(dimPos, idxMap.getContext()))) {
+              operandDimPairs.push_back({$_op->getOperand(i), *mayOperandDim});
+            }
+          }
+        }
+
+        return;
+      }]
+    >,
     //===------------------------------------------------------------------===//
     // Linalg generalization hooks.
     //===------------------------------------------------------------------===//