[mlir][Linalg] Add interface method to Linalg ops to allow setting the output operand.

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

diff --git a/mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.td b/mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.td
index c305ff0f619a29687aa52fd6ef0d16dd24ee27f2..66745096ded3481f9edb4f86df557d0eee869032 100644 (file)
--- a/mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.td
+++ b/mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.td
@@ -410,6 +410,19 @@ def LinalgStructuredInterface : OpInterface<"LinalgOp"> {
         return &this->getOperation()->getOpOperand(getNumInputs() + i);
       }]
     >,
+    InterfaceMethod<
+      /*desc=*/[{
+        Set the `i`-th output operand.
+      }],
+      /*retTy=*/"void",
+      /*methodName=*/"setOutputOperand",
+      /*args=*/(ins "int64_t":$i, "Value":$value),
+      /*methodBody=*/"",
+      /*defaultImplementation=*/[{
+        assert(i >= 0 && i < getNumOutputs());
+        this->getOperation()->setOperand(getNumInputs() + i, value);
+      }]
+    >,
     InterfaceMethod<
       /*desc=*/[{
         Return the subset of output operands that are of buffer type.