[Mlir-commits] [mlir] 91563c4 - [mlir] Modify LinalgStructuredInterface to allow the computation block to have arguments only for a subset of operands.

[Oleg Shyshkov]

Author: Oleg Shyshkov
Date: 2022-09-22T18:46:33Z
New Revision: 91563c419e57e3e22de96122247ea27540b14659

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

LOG: [mlir] Modify LinalgStructuredInterface to allow the computation block to have arguments only for a subset of operands.

Summary:
Currently there is an expectations that there is a corresponsing block argument
for each operand. For some operation, it leads to unused arguments. For example,
in `map`, only input operands are used for the computation.

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

Added: 
    

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

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.td b/mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.td
index bf2509f23beb3..64c2bd1b28510 100644
--- a/mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.td
+++ b/mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.td
@@ -570,6 +570,21 @@ def LinalgStructuredInterface : OpInterface<"LinalgOp"> {
       /*methodBody=*/"",
       /*defaultImplementation=*/""
     >,
+    InterfaceMethod<
+      /*desc=*/[{
+        Return op operands that have a corresponding argument in the basic block.
+        By default, the block should have an argument for each operand, but there
+        are expection. For example, in `map` output operand isn't used in
+        the block.
+      }],
+      /*retTy=*/"OpOperandVector",
+      /*methodName=*/"getOpOperandsMatchingBBargs",
+      /*args=*/(ins),
+      /*methodBody=*/"",
+      /*defaultImplementation=*/[{
+        return $_op.getInputAndOutputOperands();
+      }]
+    >,
     //===------------------------------------------------------------------===//
     // Linalg generalization hooks.
     //===------------------------------------------------------------------===//

diff  --git a/mlir/lib/Dialect/Linalg/IR/LinalgInterfaces.cpp b/mlir/lib/Dialect/Linalg/IR/LinalgInterfaces.cpp
index c30f3575e22af..317240089924f 100644
--- a/mlir/lib/Dialect/Linalg/IR/LinalgInterfaces.cpp
+++ b/mlir/lib/Dialect/Linalg/IR/LinalgInterfaces.cpp
@@ -762,11 +762,11 @@ LogicalResult mlir::linalg::detail::verifyStructuredOpInterface(Operation *op) {
   // not used).
   Block &block = linalgOp->getRegion(0).front();
 
-  if (linalgOp.getNumInputsAndOutputs() != block.getNumArguments())
+  if (linalgOp.getOpOperandsMatchingBBargs().size() != block.getNumArguments())
     return op->emitOpError("expected as many non-induction variable region "
                            "arguments as the number of input/output operands");
 
-  for (OpOperand *opOperand : linalgOp.getInputAndOutputOperands()) {
+  for (OpOperand *opOperand : linalgOp.getOpOperandsMatchingBBargs()) {
     Type elementType = getElementTypeOrSelf(opOperand->get());
     Type argType = block.getArgument(opOperand->getOperandNumber()).getType();
     if (elementType != argType)