[Mlir-commits] [mlir] f3ab0cc - [mlir][Linalg] Add couple of convenience methods to `LinalgInterface`.

[llvmlistbot at llvm.org]

Author: MaheshRavishankar
Date: 2022-01-25T11:46:11-08:00
New Revision: f3ab0ccd00db3bb26851d3e1cf4a789fe7fa6e2c

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

LOG: [mlir][Linalg] Add couple of convenience methods to `LinalgInterface`.

Add methods to
- Get block argument that is tied with an opOperand
- Get the yield value that is tied with a output opOperand.

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

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 413c2cc18acea..f40f82ed7cd94 100644
--- a/mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.td
+++ b/mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.td
@@ -641,6 +641,19 @@ def LinalgStructuredInterface : OpInterface<"LinalgOp"> {
         return !opOperand->get().getType().template isa<ShapedType>();
       }]
     >,
+    InterfaceMethod<
+      /*desc=*/[{
+        Return the block argument for an `opOperand`.
+      }],
+      /*retTy=*/"BlockArgument",
+      /*methodName=*/"getTiedBlockArgument",
+      /*args=*/(ins "OpOperand *":$opOperand),
+      /*methodBody=*/"",
+      /*defaultImplementation=*/[{
+        assert(opOperand->getOwner() == this->getOperation());
+        return getBlock()->getArgument(opOperand->getOperandNumber());
+      }]
+    >,
     InterfaceMethod<
       /*desc=*/[{
         Return the input or output indexing map for `opOperand`.
@@ -672,6 +685,24 @@ def LinalgStructuredInterface : OpInterface<"LinalgOp"> {
         return this->getOperation()->getResult(resultIndex);
       }]
     >,
+    InterfaceMethod<
+      /*desc=*/[{
+        Return the value yielded by the region corresponding to an output
+        `opOperand`.
+      }],
+      /*retTy=*/"OpOperand *",
+      /*methodName=*/"getTiedYieldValue",
+      /*args=*/(ins "OpOperand*":$opOperand),
+      /*methodBody=*/"",
+      /*defaultImplementation=*/[{
+        assert(opOperand->getOwner() == this->getOperation());
+        int64_t resultIndex = opOperand->getOperandNumber() - getNumInputs();
+        assert(resultIndex >= 0 &&
+               resultIndex < this->getOperation()->getNumResults());
+        Operation *yieldOp = getBlock()->getTerminator();
+        return &yieldOp->getOpOperand(resultIndex);
+      }]
+    >,
     //===------------------------------------------------------------------===//
     // Other interface methods.
     //===------------------------------------------------------------------===//