[mlir][scf] Canonicalize scf.while with unused results

[bondhugula]

mlir/lib/Dialect/SCF/SCF.cpp

2259

Doc comment please.

2268–2315

Code comments needed for the major blocks.

2306

Stray semi-colon?

mlir/test/Dialect/SCF/canonicalize.mlir

822–830

%arg0 <--- avoid matching numbered SSA values. Use .* or match.

[Hardcode84]

review comments

[bondhugula]

mlir/lib/Dialect/SCF/SCF.cpp

2259

Nit: Consider rephrasing this way to avoid ambiguity.

/// Remove WhileOp results that are also unused in after block.

2259

Terminate comments with a full stop -- here and everywhere below.

2318

Code comment here.

[Hardcode84]

mlir/lib/Dialect/SCF/SCF.cpp

2318

In this specific block everything should clear from the code itself, commenting it like // Replacing old condition terminator will just add more noise

[Hardcode84]

fix comments

[ftynse]

mlir/lib/Dialect/SCF/SCF.cpp

2326

In-place mutation of block argument may be problematic if this ever used in the dialect conversion infra (greedy rewriter in the canonicalization is fine though). Please add a big warning in the top-level documentation of this pattern and a remark comment here.

Or, ideally, figure out a way of achieving the same behavior. Something like mergeBlocks could work.

[Hardcode84]

use mergeBlocks

mlir/lib/Dialect/SCF/SCF.cpp

for (auto yieldedAndBlockArgs :
           std::get<1>(yieldedAndBlockArgs).replaceAllUsesWith(constantTrue);
           replaced = true;
         }
       }
     }
     return success(replaced);
   }
 };
+
+/// Remove WhileOp results that are also unused in 'after' block.
+///
+///  %0:2 = scf.while () : () -> (i32, i64) {
+///    %condition = "test.condition"() : () -> i1
+///    %v1 = "test.get_some_value"() : () -> i32
+///    %v2 = "test.get_some_value"() : () -> i64
+///    scf.condition(%condition) %v1, %v2 : i32, i64
+///  } do {
+///  ^bb0(%arg0: i32, %arg1: i64):
+///    "test.use"(%arg0) : (i32) -> ()
+///    scf.yield
+///  }
+///  return %0#0 : i32
+///
+/// becomes
+///  %0 = scf.while () : () -> (i32) {
+///    %condition = "test.condition"() : () -> i1
+///    %v1 = "test.get_some_value"() : () -> i32
+///    %v2 = "test.get_some_value"() : () -> i64
+///    scf.condition(%condition) %v1 : i32
+///  } do {
+///  ^bb0(%arg0: i32):
+///    "test.use"(%arg0) : (i32) -> ()
+///    scf.yield
+///  }
+///  return %0 : i32
+struct WhileUnusedResult : public OpRewritePattern<WhileOp> {
+  using OpRewritePattern<WhileOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(WhileOp op,
+                                PatternRewriter &rewriter) const override {
+    auto term = op.getConditionOp();
+    auto afterArgs = op.getAfterArguments();
+    auto termArgs = term.args();
+
+    // Collect results mapping, new terminator args and new result types.
+    SmallVector<unsigned> newResultsIndices;
+    SmallVector<Type> newResultTypes;
+    SmallVector<Value> newTermArgs;
+    bool needUpdate = false;
+    for (auto it :
+         llvm::enumerate(llvm::zip(op.getResults(), afterArgs, termArgs))) {
+      auto i = static_cast<unsigned>(it.index());
+      Value result = std::get<0>(it.value());
+      Value afterArg = std::get<1>(it.value());
+      Value termArg = std::get<2>(it.value());
+      if (result.use_empty() && afterArg.use_empty()) {
+        needUpdate = true;
+      } else {
+        newResultsIndices.emplace_back(i);
+        newTermArgs.emplace_back(termArg);
+        newResultTypes.emplace_back(result.getType());
+      }
+    }
+
+    if (!needUpdate)
+      return failure();
+
+    {
+      OpBuilder::InsertionGuard g(rewriter);
+      rewriter.setInsertionPoint(term);
+      rewriter.replaceOpWithNewOp<ConditionOp>(term, term.condition(),
+                                               newTermArgs);
+    }
+
+    auto newWhile =
+        rewriter.create<WhileOp>(op.getLoc(), newResultTypes, op.inits());
+
+    Block &newAfterBlock = *rewriter.createBlock(
+        &newWhile.after(), /*insertPt*/ {}, newResultTypes);
+
+    // Build new results list and new after block args (unused entries will be
+    // null).
+    SmallVector<Value> newResults(op.getNumResults());
+    SmallVector<Value> newAfterBlockArgs(op.getNumResults());
+    for (auto it : llvm::enumerate(newResultsIndices)) {
+      newResults[it.value()] = newWhile.getResult(it.index());
+      newAfterBlockArgs[it.value()] = newAfterBlock.getArgument(it.index());
+    }
+
+    rewriter.inlineRegionBefore(op.before(), newWhile.before(),
+                                newWhile.before().begin());
+
+    Block &afterBlock = op.after().front();
+    rewriter.mergeBlocks(&afterBlock, &newAfterBlock, newAfterBlockArgs);
+
+    rewriter.replaceOp(op, newResults);
+    return success();
+  }
+};
 } // namespace
 
 void WhileOp::getCanonicalizationPatterns(OwningRewritePatternList &results,
                                           MLIRContext *context) {
-  results.insert<WhileConditionTruth>(context);
+  results.insert<WhileConditionTruth, WhileUnusedResult>(context);
 }
 
 //===----------------------------------------------------------------------===//
 // TableGen'd op method definitions
 //===----------------------------------------------------------------------===//
 
 #define GET_OP_CLASSES
 #include "mlir/Dialect/SCF/SCFOps.cpp.inc"

mlir/test/Dialect/SCF/canonicalize.mlir

 // CHECK-NEXT:       %[[sv2:.+]] = "test.get_some_value"() : () -> i32
 // CHECK-NEXT:       scf.yield %[[sv2]] : i32
 // CHECK-NEXT:     }
 // CHECK-NEXT:     return %[[if]], %arg1 : i32, i64
 
 // -----
 
 // CHECK-LABEL: @while_cond_true
-func @while_cond_true() {
+func @while_cond_true() -> i1 {
   %0 = scf.while () : () -> i1 {
     %condition = "test.condition"() : () -> i1
     scf.condition(%condition) %condition : i1
   } do {
   ^bb0(%arg0: i1):
     "test.use"(%arg0) : (i1) -> ()
     scf.yield
   }
-  return
+  return %0 : i1
 }
 // CHECK-NEXT:         %[[true:.+]] = arith.constant true
 // CHECK-NEXT:         %{{.+}} = scf.while : () -> i1 {
 // CHECK-NEXT:           %[[cmp:.+]] = "test.condition"() : () -> i1
 // CHECK-NEXT:           scf.condition(%[[cmp]]) %[[cmp]] : i1
 // CHECK-NEXT:         } do {
 // CHECK-NEXT:         ^bb0(%arg0: i1):  // no predecessors
 // CHECK-NEXT:           "test.use"(%[[true]]) : (i1) -> ()
 // CHECK-NEXT:           scf.yield
 // CHECK-NEXT:         }
 
 // -----
 
+// CHECK-LABEL: @while_unused_result
+func @while_unused_result() -> i32 {
+  %0:2 = scf.while () : () -> (i32, i64) {
+    %condition = "test.condition"() : () -> i1
+    %v1 = "test.get_some_value"() : () -> i32
+    %v2 = "test.get_some_value"() : () -> i64
+    scf.condition(%condition) %v1, %v2 : i32, i64
+  } do {
+  ^bb0(%arg0: i32, %arg1: i64):
+    "test.use"(%arg0) : (i32) -> ()
+    scf.yield
+  }
+  return %0#0 : i32
+}
+// CHECK-NEXT:         %[[res:.*]] = scf.while : () -> i32 {
+// CHECK-NEXT:           %[[cmp:.*]] = "test.condition"() : () -> i1
+// CHECK-NEXT:           %[[val:.*]] = "test.get_some_value"() : () -> i32
+// CHECK-NEXT:           %{{.*}} = "test.get_some_value"() : () -> i64
+// CHECK-NEXT:           scf.condition(%[[cmp]]) %[[val]] : i32
+// CHECK-NEXT:         } do {
+// CHECK-NEXT:         ^bb0(%[[arg:.*]]: i32):  // no predecessors
+// CHECK-NEXT:           "test.use"(%[[arg]]) : (i32) -> ()
+// CHECK-NEXT:           scf.yield
+// CHECK-NEXT:         }
+// CHECK-NEXT:         return %[[res]] : i32
+
+// -----
+
 // CHECK-LABEL: @combineIfs
 func @combineIfs(%arg0 : i1, %arg2: i64) -> (i32, i32) {
   %res = scf.if %arg0 -> i32 {
     %v = "test.firstCodeTrue"() : () -> i32
     scf.yield %v : i32
   } else {
     %v2 = "test.firstCodeFalse"() : () -> i32
     scf.yield %v2 : i32