Update. Related test: test/Dialect/SCF/transform-ops.mlir

Author: lialan

mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td

@@ -1475,7 +1475,7 @@ def Tensor_ParallelInsertSliceOp : Tensor_Op<"parallel_insert_slice", [
        OffsetSizeAndStrideOpInterface,
        DeclareOpInterfaceMethods<ParallelCombiningOpInterface,
           ["getUpdatedDestinations", "getIteratingParent",
-           "promoteInParallelLoop"]>,
+           "promoteInParallelLoop", "canPromoteInParallelLoop"]>,
        // TODO: Cannot use an interface here atm, verify this manually for now.
        // HasParent<"InParallelOpInterface">
   ]> {

mlir/include/mlir/Interfaces/ParallelCombiningOpInterface.h

@@ -17,7 +17,6 @@
 #include "mlir/IR/OpDefinition.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/Support/LogicalResult.h"
-#include "llvm/ADT/SmallVector.h"
 
 namespace mlir {
 namespace detail {

mlir/include/mlir/Interfaces/ParallelCombiningOpInterface.td

@@ -109,10 +109,10 @@ def ParallelCombiningOpInterface : OpInterface<"ParallelCombiningOpInterface"> {
     InterfaceMethod<
       /*desc=*/[{
         Promotes this parallel combining op out of its enclosing parallel loop
-        and returns the values that should replace the destinations updated by
+        and returns the value that should replace the destination updated by
         this op.
       }],
-      /*retTy=*/"::mlir::FailureOr<::llvm::SmallVector<::mlir::Value>>",
+      /*retTy=*/"::mlir::FailureOr<::mlir::Value>",
       /*methodName=*/"promoteInParallelLoop",
       /*args=*/(ins "::mlir::RewriterBase &":$rewriter)
     >,
@@ -124,7 +124,8 @@ def ParallelCombiningOpInterface : OpInterface<"ParallelCombiningOpInterface"> {
       /*retTy=*/"bool",
       /*methodName=*/"canPromoteInParallelLoop",
       /*args=*/(ins "::mlir::RewriterBase &":$rewriter),
-      /*methodBody=*/[{ return true; }]
+      /*methodBody=*/"",
+      /*defaultImplementation=*/[{ return false; }]
     >,
   ];
 }

mlir/lib/Dialect/SCF/IR/SCF.cpp

@@ -672,14 +672,12 @@ LogicalResult mlir::scf::promote(RewriterBase &rewriter, scf::ForallOp forallOp)
     auto parallelCombiningOp =
         dyn_cast<ParallelCombiningOpInterface>(&yieldingOp);
     if (!parallelCombiningOp)
-      return rewriter.notifyMatchFailure(
-          forallOp, "terminator has non-parallel-combining op");
+      continue;
     if (!parallelCombiningOp.canPromoteInParallelLoop(rewriter))
       return rewriter.notifyMatchFailure(
           forallOp, "parallel combining op cannot be promoted");
   }
 
-
   // Replace block arguments with lower bounds (replacements for IVs) and
   // outputs.
   SmallVector<Value> bbArgReplacements = forallOp.getLowerBound(rewriter);
@@ -702,12 +700,12 @@ LogicalResult mlir::scf::promote(RewriterBase &rewriter, scf::ForallOp forallOp)
 
     assert(parallelCombiningOp.canPromoteInParallelLoop(rewriter));
 
-    FailureOr<SmallVector<Value>> promotedValues =
+    FailureOr<Value> promotedValue =
         parallelCombiningOp.promoteInParallelLoop(rewriter);
-    if (failed(promotedValues))
+    if (failed(promotedValue))
       return failure();
 
-    results.append(promotedValues->begin(), promotedValues->end());
+    results.push_back(*promotedValue);
   }
   if (results.size() != forallOp.getResults().size())
     return rewriter.notifyMatchFailure(

mlir/lib/Dialect/Tensor/IR/TensorOps.cpp

@@ -3947,19 +3947,23 @@ Operation *ParallelInsertSliceOp::getIteratingParent() {
   return nullptr;
 }
 
-FailureOr<SmallVector<Value>>
+FailureOr<Value>
 ParallelInsertSliceOp::promoteInParallelLoop(RewriterBase &rewriter) {
   Value dst = getDest();
   Value src = getSource();
   if (!isa<TensorType>(src.getType()))
-    return rewriter.notifyMatchFailure(getOperation(),
-                                       "expected tensor source");
+    return failure();
 
   Value inserted = tensor::InsertSliceOp::create(
       rewriter, getLoc(), dst.getType(), src, dst, getOffsets(), getSizes(),
       getStrides(), getStaticOffsets(), getStaticSizes(), getStaticStrides());
 
-  return SmallVector<Value>{inserted};
+  return inserted;
+}
+
+bool ParallelInsertSliceOp::canPromoteInParallelLoop(RewriterBase &) {
+  return isa<TensorType>(getSource().getType()) &&
+         isa<TensorType>(getDest().getType());
 }
 
 //===----------------------------------------------------------------------===//