Code review comments

Author: ddubov100

mlir/include/mlir/Dialect/SCF/IR/SCFOps.td

@@ -78,8 +78,7 @@ def ConditionOp : SCF_Op<"condition", [
 
 def ExecuteRegionOp : SCF_Op<"execute_region", [
     DeclareOpInterfaceMethods<RegionBranchOpInterface>,
-    RecursiveMemoryEffects,
-    DeclareOpInterfaceMethods<MemoryEffectsOpInterface>]> {//, SingleBlockImplicitTerminator<"scf::YieldOp">]> { //, RecursiveMemoryEffects]> {
+    RecursiveMemoryEffects]> {
   let summary = "operation that executes its region exactly once";
   let description = [{
     The `scf.execute_region` operation is used to allow multiple blocks within SCF

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

@@ -27,6 +27,7 @@
 #include "mlir/Interfaces/ValueBoundsOpInterface.h"
 #include "mlir/Transforms/InliningUtils.h"
 #include "llvm/ADT/MapVector.h"
+#include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/Support/Casting.h"
 #include "llvm/Support/DebugLog.h"
@@ -306,20 +307,25 @@ struct ExecuteRegionForwardingEliminator
     SmallVector<Operation *> yieldOps;
     for (Block &block : op.getRegion()) {
       if (auto yield = dyn_cast<scf::YieldOp>(block.getTerminator())) {
-        if (yield.getResults().empty())
+        if (yield.getOperands().empty())
           continue;
         yieldOps.push_back(yield.getOperation());
       }
     }
 
-    if (yieldOps.size() != 1)
+    if (yieldOps.empty())
       return failure();
 
-    auto yieldOp = cast<scf::YieldOp>(yieldOps.front());
-    auto yieldedValues = yieldOp.getOperands();
+    // Check if all yield operations have the same operands.
+    auto yieldOpsOperands = yieldOps[0]->getOperands();
+    for (auto *yieldOp : yieldOps) {
+      if (yieldOp->getOperands() != yieldOpsOperands)
+        return failure();
+    }
+
     // Check if all yielded values are from outside the region
     bool allExternal = true;
-    for (Value yieldedValue : yieldedValues) {
+    for (Value yieldedValue : yieldOpsOperands) {
       if (isValueFromInsideRegion(yieldedValue, op)) {
         allExternal = false;
         break;
@@ -337,13 +343,16 @@ struct ExecuteRegionForwardingEliminator
     // Move the region content to the new operation
     newOp.getRegion().takeBody(op.getRegion());
 
-    // Replace the yield operation with a new yield operation with no results.
-    rewriter.setInsertionPoint(yieldOp);
-    rewriter.eraseOp(yieldOp);
-    rewriter.create<scf::YieldOp>(yieldOp.getLoc());
+    // Replace all yield operations with a new yield operation with no results.
+    // scf.execute_region must have at least one yield operation.
+    for (auto *yieldOp : yieldOps) {
+      rewriter.setInsertionPoint(yieldOp);
+      rewriter.eraseOp(yieldOp);
+      rewriter.create<scf::YieldOp>(yieldOp->getLoc());
+    }
 
     // Replace the old operation with the external values directly.
-    rewriter.replaceOp(op, yieldedValues);
+    rewriter.replaceOp(op, yieldOpsOperands);
     return success();
   }
 
@@ -352,12 +361,11 @@ struct ExecuteRegionForwardingEliminator
                                ExecuteRegionOp executeRegionOp) const {
     // Check if the value is defined within the execute_region
     if (Operation *defOp = value.getDefiningOp())
-      return executeRegionOp.getRegion() = defOp->getParentRegion();
+      return &executeRegionOp.getRegion() == defOp->getParentRegion();
 
     // If it's a block argument, check if it's from within the region
-    if (BlockArgument blockArg = dyn_cast<BlockArgument>(value)) {
-      return executeRegionOp.getRegion() == blockArg.getParentRegion());
-    }
+    if (BlockArgument blockArg = dyn_cast<BlockArgument>(value))
+      return &executeRegionOp.getRegion() == blockArg.getParentRegion();
 
     return false; // Value is from outside the region
   }
@@ -369,17 +377,6 @@ void ExecuteRegionOp::getCanonicalizationPatterns(RewritePatternSet &results,
               ExecuteRegionForwardingEliminator>(context);
 }
 
-void ExecuteRegionOp::getEffects(
-    SmallVectorImpl<SideEffects::EffectInstance<MemoryEffects::Effect>>
-        &effects) {
-  if (!getNoInline())
-    return;
-  // In case there is attribute no_inline we want the region not to be inlined
-  // into the parent operation.
-  effects.emplace_back(MemoryEffects::Write::get(),
-                       SideEffects::DefaultResource::get());
-}
-
 void ExecuteRegionOp::getSuccessorRegions(
     RegionBranchPoint point, SmallVectorImpl<RegionSuccessor> &regions) {
   // If the predecessor is the ExecuteRegionOp, branch into the body.

mlir/test/Dialect/Bufferization/Transforms/one-shot-module-bufferize.mlir

@@ -377,21 +377,6 @@ func.func private @execute_region_test(%t1 : tensor<?xf32>)
   // CHECK: return %{{.*}}, %{{.*}} : f32, f32
   return %0, %1, %2 : f32, tensor<?xf32>, f32
 }
-
-// -----
-
-// CHECK-LABEL: func @no_inline_execute_region_not_canonicalized
-func.func @no_inline_execute_region_not_canonicalized() {
-  %c = arith.constant 42 : i32
-  // CHECK: scf.execute_region
-  // CHECK-SAME: no_inline
-  %v = scf.execute_region -> i32 no_inline {
-    scf.yield %c : i32
-  }
-  // CHECK: return
-  return
-}
-
 // -----
 
 //      CHECK:  func private @some_external_func(memref<?xf32, strided<[?], offset: ?>>)

mlir/test/Dialect/SCF/canonicalize.mlir

@@ -1604,24 +1604,141 @@ func.func @func_execute_region_inline_multi_yield() {
 
 // -----
 
-// CHECK-LABEL:   func.func private @canonicalize_execute_region_yeilding_external_value(
-// CHECK-SAME:                                                                           %[[VAL_0:.*]]: tensor<1x120xui8>) -> tensor<1x60xui8> {
-// CHECK:           %[[VAL_1:.*]] = memref.alloc() {alignment = 64 : i64} : memref<1x60xui8>
+// Test case with single scf.yield op inside execute_region and its operand is defined outside the execute_region op.
+// Make scf.execute_region not to return anything.
+
 // CHECK:           scf.execute_region no_inline {
 // CHECK:             scf.yield
 // CHECK:           }
-// CHECK:           %[[VAL_2:.*]] = bufferization.to_tensor %[[VAL_1]] : memref<1x60xui8> to tensor<1x60xui8>
-// CHECK:           return %[[VAL_2]] : tensor<1x60xui8>
-// CHECK:         }
 
-func.func private @canonicalize_execute_region_yeilding_external_value(%arg0: tensor<1x120xui8>) -> tensor<1x60xui8> {
-  %alloc = memref.alloc() {alignment = 64 : i64} : memref<1x60xui8>
-  %0 = bufferization.to_tensor %alloc : memref<1x60xui8> to tensor<1x60xui8>
-  %1 = scf.execute_region -> memref<1x60xui8> no_inline {
-    scf.yield %alloc : memref<1x60xui8>
+module {
+func.func private @foo()->()
+func.func private @execute_region_yeilding_external_value() -> memref<1x60xui8> {
+  %alloc = memref.alloc() {alignment = 64 : i64} : memref<1x60xui8>  
+  %1 = scf.execute_region -> memref<1x60xui8> no_inline {    
+    func.call @foo():()->()
+    scf.yield %alloc: memref<1x60xui8>
+  }  
+  return %1 : memref<1x60xui8>
+}
+}
+
+// -----
+
+// Test case with scf.yield op inside execute_region with multiple operands.
+// One of operands is defined outside the execute_region op.
+// In this case scf.execute_region doesn't change.
+
+// CHECK:           %[[VAL_1:.*]]:2 = scf.execute_region -> (memref<1x60xui8>, memref<1x120xui8>) no_inline {
+// CHECK:             scf.yield %{{.*}}, %{{.*}} : memref<1x60xui8>, memref<1x120xui8>
+
+module {
+func.func private @foo()->()
+func.func private @execute_region_yeilding_external_and_local_values() -> (memref<1x60xui8>, memref<1x120xui8>) {
+  %alloc = memref.alloc() {alignment = 64 : i64} : memref<1x60xui8>  
+  %1, %2 = scf.execute_region -> (memref<1x60xui8>, memref<1x120xui8>) no_inline {    
+    %alloc_1 = memref.alloc() {alignment = 64 : i64} : memref<1x120xui8>
+    func.call @foo():()->()
+    scf.yield %alloc, %alloc_1: memref<1x60xui8>,  memref<1x120xui8>
+  }  
+  return %1, %2 : memref<1x60xui8>, memref<1x120xui8>
+}
+}
+
+// -----
+
+// Test case with multiple scf.yield ops inside execute_region with same operands and those operands are defined outside the execute_region op..
+// Make scf.execute_region not to return anything.
+// scf.yield must remain, cause scf.execute_region can't be empty.
+
+// CHECK:           scf.execute_region no_inline {
+// CHECK:             %[[VAL_3:.*]] = "test.cmp"() : () -> i1
+// CHECK:             cf.cond_br %[[VAL_3]], ^bb1, ^bb2
+// CHECK:           ^bb1:
+// CHECK:             scf.yield
+// CHECK:           ^bb2:
+// CHECK:             scf.yield
+// CHECK:           }
+
+module {
+  func.func private @foo()->()
+  func.func private @execute_region_multiple_yields_same_operands() -> (memref<1x60xui8>, memref<1x120xui8>) {
+    %alloc = memref.alloc() {alignment = 64 : i64} : memref<1x60xui8>  
+    %alloc_1 = memref.alloc() {alignment = 64 : i64} : memref<1x120xui8>  
+    %1, %2 = scf.execute_region -> (memref<1x60xui8>, memref<1x120xui8>) no_inline {
+      %c = "test.cmp"() : () -> i1
+      cf.cond_br %c, ^bb2, ^bb3
+    ^bb2:    
+      func.call @foo():()->()
+      scf.yield %alloc, %alloc_1 : memref<1x60xui8>, memref<1x120xui8>
+    ^bb3: 
+      func.call @foo():()->()   
+      scf.yield %alloc, %alloc_1 : memref<1x60xui8>, memref<1x120xui8>
+    }  
+    return %1, %2 : memref<1x60xui8>, memref<1x120xui8>
   }
-  %2 = bufferization.to_tensor %1 : memref<1x60xui8> to tensor<1x60xui8>
-  return %2 : tensor<1x60xui8>
+}
+
+// -----
+
+// Test case with multiple scf.yield ops with at least one different operand, then no change.
+
+// CHECK:           %[[VAL_3:.*]]:2 = scf.execute_region -> (memref<1x60xui8>, memref<1x120xui8>) no_inline {
+// CHECK:           ^bb1:
+// CHECK:             scf.yield %{{.*}}, %{{.*}} : memref<1x60xui8>, memref<1x120xui8>
+// CHECK:           ^bb2:
+// CHECK:             scf.yield %{{.*}}, %{{.*}} : memref<1x60xui8>, memref<1x120xui8>
+// CHECK:           }
+
+module {
+  func.func private @foo()->()
+  func.func private @execute_region_multiple_yields_different_operands() -> (memref<1x60xui8>, memref<1x120xui8>) {
+    %alloc = memref.alloc() {alignment = 64 : i64} : memref<1x60xui8>  
+    %alloc_1 = memref.alloc() {alignment = 64 : i64} : memref<1x120xui8>  
+    %alloc_2 = memref.alloc() {alignment = 64 : i64} : memref<1x120xui8>  
+    %1, %2 = scf.execute_region -> (memref<1x60xui8>, memref<1x120xui8>) no_inline {
+      %c = "test.cmp"() : () -> i1
+      cf.cond_br %c, ^bb2, ^bb3
+    ^bb2:    
+      func.call @foo():()->()
+      scf.yield %alloc, %alloc_1 : memref<1x60xui8>, memref<1x120xui8>
+    ^bb3: 
+      func.call @foo():()->()   
+      scf.yield %alloc, %alloc_2 : memref<1x60xui8>, memref<1x120xui8>
+    }  
+    return %1, %2 : memref<1x60xui8>, memref<1x120xui8>
+  }
+}
+
+// -----
+
+// Test case with multiple scf.yield ops each has different operand.
+// In this case scf.execute_region isn't changed.
+
+// CHECK:           %[[VAL_2:.*]] = scf.execute_region -> memref<1x60xui8> no_inline {
+// CHECK:           ^bb1:
+// CHECK:             scf.yield %{{.*}} : memref<1x60xui8>
+// CHECK:           ^bb2:
+// CHECK:             scf.yield %{{.*}} : memref<1x60xui8>
+// CHECK:           }
+
+module {
+func.func private @foo()->()
+func.func private @execute_region_multiple_yields_different_operands() -> (memref<1x60xui8>) {
+  %alloc = memref.alloc() {alignment = 64 : i64} : memref<1x60xui8>  
+  %alloc_1 = memref.alloc() {alignment = 64 : i64} : memref<1x60xui8>   
+  %1 = scf.execute_region -> (memref<1x60xui8>) no_inline {
+    %c = "test.cmp"() : () -> i1
+    cf.cond_br %c, ^bb2, ^bb3
+  ^bb2:    
+    func.call @foo():()->()
+    scf.yield %alloc : memref<1x60xui8>
+  ^bb3:    
+    func.call @foo():()->()
+    scf.yield %alloc_1 : memref<1x60xui8>
+  }    
+  return %1 : memref<1x60xui8>
+}
 }
 
 // -----