[Mlir-commits] [mlir] 5b66e44 - [mlir][SCF] Make the scf.take_assumed_branch transform only read its target handle

[Nicolas Vasilache]

Author: Nicolas Vasilache
Date: 2023-04-12T12:30:18-07:00
New Revision: 5b66e44cff1ff19fc72a264b00e8b83250c86863

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

LOG: [mlir][SCF] Make the scf.take_assumed_branch transform only read its target handle

Handles are tracked properly and usage composes better if we don't artificially lose handles.

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

Added: 
    

Modified: 
    mlir/lib/Dialect/SCF/TransformOps/SCFTransformOps.cpp
    mlir/test/Dialect/SCF/transform-op-take-assumed-branch.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/SCF/TransformOps/SCFTransformOps.cpp b/mlir/lib/Dialect/SCF/TransformOps/SCFTransformOps.cpp
index b87fc777cdc4b..f1b45e1f09ace 100644
--- a/mlir/lib/Dialect/SCF/TransformOps/SCFTransformOps.cpp
+++ b/mlir/lib/Dialect/SCF/TransformOps/SCFTransformOps.cpp
@@ -282,7 +282,7 @@ DiagnosedSilenceableFailure transform::TakeAssumedBranchOp::applyToOne(
 
 void transform::TakeAssumedBranchOp::getEffects(
     SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
-  consumesHandle(getTarget(), effects);
+  onlyReadsHandle(getTarget(), effects);
   modifiesPayload(effects);
 }
 

diff  --git a/mlir/test/Dialect/SCF/transform-op-take-assumed-branch.mlir b/mlir/test/Dialect/SCF/transform-op-take-assumed-branch.mlir
index 15d9e56ad099a..68325702ac282 100644
--- a/mlir/test/Dialect/SCF/transform-op-take-assumed-branch.mlir
+++ b/mlir/test/Dialect/SCF/transform-op-take-assumed-branch.mlir
@@ -12,7 +12,6 @@ transform.sequence failures(propagate) {
 ^bb0(%arg1: !transform.any_op):
   %if = transform.structured.match ops{["scf.if"]} in %arg1 
     : (!transform.any_op) -> !transform.any_op
-
   // expected-error @+1 {{requires an scf.if op with a single-block `else` region}}
   transform.scf.take_assumed_branch %if take_else_branch 
     : (!transform.any_op) -> ()
@@ -20,6 +19,30 @@ transform.sequence failures(propagate) {
 
 // -----
 
+// CHECK-LABEL: if_no_else
+func.func @if_no_else(%cond: i1, %a: index, %b: memref<?xf32>, %c: i8) {
+  scf.if %cond {
+    "some_op"(%cond, %b) : (i1, memref<?xf32>) -> ()
+    scf.yield
+  }
+  return
+}
+
+transform.sequence failures(propagate) {
+^bb0(%arg1: !transform.any_op):
+  %if = transform.structured.match ops{["scf.if"]} in %arg1 
+    : (!transform.any_op) -> !transform.any_op
+  %some_op = transform.structured.match ops{["some_op"]} in %arg1 
+    : (!transform.any_op) -> !transform.any_op
+
+  transform.scf.take_assumed_branch %if : (!transform.any_op) -> ()
+  
+  // Handle to formerly nested `some_op` is still valid after the transform.
+  transform.print %some_op: !transform.any_op
+}
+
+// -----
+
 // CHECK-LABEL: tile_tensor_pad
 func.func @tile_tensor_pad(
   %arg0 : tensor<?x?xf32>, %cst : f32, %low: index, %high: index)