From e2f1d5ce9147118c94e2d65ea5c53e7474e8ff4a Mon Sep 17 00:00:00 2001 From: Diego Caballero Date: Wed, 29 Mar 2023 19:09:41 +0000 Subject: [PATCH] [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 --- .../mlir/Dialect/Linalg/IR/LinalgInterfaces.td | 25 ++++++++++++++++++++++ 1 file changed, 25 insertions(+) diff --git a/mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.td b/mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.td index 9b51228..ace2427 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>&":$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. //===------------------------------------------------------------------===// -- 2.7.4