[mlir] add scf.if op canonicalization pattern that removes unused results

[gysit]

Summary

The patch adds a canonicalization pattern that removes the unused results of scf.if operation. As a result, cse may remove unused computations in the then and else regions of the scf.if operation.

[mehdi_amini]

mlir/lib/Dialect/SCF/SCF.cpp

512

Nit: if op -> IfOp

529

Note: operation can have their operands erased in place, seems like it would be possible here.

540

Please early exit: if (usedResults.size() != op.getNumResults()) return failure()

[gysit]

• exit early if the pattern does not apply
• update terminators of then and else bodies in place using the setOperands method

[gysit]

mlir/lib/Dialect/SCF/SCF.cpp

529

Thanks for the hint! I was not aware that the pattern rewriter has this in place update functionality.

[gysit]

update to latest master

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

OpBuilder getThenBodyBuilder(OpBuilder::Listener *listener = nullptr) {
                                : OpBuilder::atBlockEnd(body, listener);
     }
     OpBuilder getElseBodyBuilder(OpBuilder::Listener *listener = nullptr) {
       Block* body = getBody(1);
       return results().empty() ? OpBuilder::atBlockTerminator(body, listener)
                                : OpBuilder::atBlockEnd(body, listener);
     }
   }];
+
+  let hasCanonicalizer = 1;
 }
 
 def ParallelOp : SCF_Op<"parallel",
     [AttrSizedOperandSegments,
      DeclareOpInterfaceMethods<LoopLikeOpInterface>,
      RecursiveSideEffects,
      SingleBlockImplicitTerminator<"scf::YieldOp">]> {
   let summary = "parallel for operation";

mlir/lib/Dialect/SCF/SCF.cpp

if (auto condAttr = operands.front().dyn_cast_or_null<IntegerAttr>()) {
     regions.push_back(RegionSuccessor(elseRegion));
     return;
   }
 
   // Add the successor regions using the condition.
   regions.push_back(RegionSuccessor(condition ? &thenRegion() : elseRegion));
 }
 
+namespace {
+// Pattern to remove unused IfOp results.
+struct RemoveUnusedResults : public OpRewritePattern<IfOp> {
+  using OpRewritePattern<IfOp>::OpRewritePattern;
+
+  void transferBody(Block *source, Block *dest, ArrayRef<OpResult> usedResults,
+                    PatternRewriter &rewriter) const {
+    // Move all operations to the destination block.
+    rewriter.mergeBlocks(source, dest);
+    // Replace the yield op by one that returns only the used values.
+    auto yieldOp = cast<scf::YieldOp>(dest->getTerminator());
+    SmallVector<Value, 4> usedOperands;
+    llvm::transform(usedResults, std::back_inserter(usedOperands),
+                    [&](OpResult result) {
+                      return yieldOp.getOperand(result.getResultNumber());
+                    });
+    rewriter.updateRootInPlace(
+        yieldOp, [&]() { yieldOp.getOperation()->setOperands(usedOperands); });
+  }
+
+  LogicalResult matchAndRewrite(IfOp op,
+                                PatternRewriter &rewriter) const override {
+    // Compute the list of used results.
+    SmallVector<OpResult, 4> usedResults;
+    llvm::copy_if(op.getResults(), std::back_inserter(usedResults),
+                  [](OpResult result) { return !result.use_empty(); });
+
+    // Replace the operation if only a subset of its results have uses.
+    if (usedResults.size() == op.getNumResults())
+      return failure();
+
+    // Compute the result types of the replacement operation.
+    SmallVector<Type, 4> newTypes;
+    llvm::transform(usedResults, std::back_inserter(newTypes),
+                    [](OpResult result) { return result.getType(); });
+
+    // Create a replacement operation with empty then and else regions.
+    auto emptyBuilder = [](OpBuilder &, Location) {};
+    auto newOp = rewriter.create<IfOp>(op.getLoc(), newTypes, op.condition(),
+                                       emptyBuilder, emptyBuilder);
+
+    // Move the bodies and replace the terminators (note there is a then and
+    // an else region since the operation returns results).
+    transferBody(op.getBody(0), newOp.getBody(0), usedResults, rewriter);
+    transferBody(op.getBody(1), newOp.getBody(1), usedResults, rewriter);
+
+    // Replace the operation by the new one.
+    SmallVector<Value, 4> repResults(op.getNumResults());
+    for (auto en : llvm::enumerate(usedResults))
+      repResults[en.value().getResultNumber()] = newOp.getResult(en.index());
+    rewriter.replaceOp(op, repResults);
+    return success();
+  }
+};
+} // namespace
+
+void IfOp::getCanonicalizationPatterns(OwningRewritePatternList &results,
+                                       MLIRContext *context) {
+  results.insert<RemoveUnusedResults>(context);
+}
+
 //===----------------------------------------------------------------------===//
 // ParallelOp
 //===----------------------------------------------------------------------===//
 
 void ParallelOp::build(
     OpBuilder &builder, OperationState &result, ValueRange lowerBounds,
     ValueRange upperBounds, ValueRange steps, ValueRange initVals,
     function_ref<void(OpBuilder &, Location, ValueRange, ValueRange)>

mlir/test/Dialect/SCF/canonicalize.mlir

 // CHECK:           [[C0:%.*]] = constant 0 : index
 // CHECK:           [[C1:%.*]] = constant 1 : index
 // CHECK:           [[C42:%.*]] = constant 42 : i32
 // CHECK:           scf.parallel ([[V1:%.*]]) = ([[C0]]) to ([[C0]]) step ([[C1]]) {
 // CHECK:             store [[C42]], [[ARG0]]{{\[}}[[C0]], [[V1]]] : memref<?x?xi32>
 // CHECK:             scf.yield
 // CHECK:           }
 // CHECK:           return
+
+// -----
+
+func @one_unused() -> (index) {
+  %c0 = constant 0 : index
+  %c1 = constant 1 : index
+  %true = constant true
+  %0, %1 = scf.if %true -> (index, index) {
+    scf.yield %c0, %c1 : index, index
+  } else {
+    scf.yield %c0, %c1 : index, index
+  }
+  return %1 : index
+}
+
+// CHECK-LABEL:   func @one_unused
+// CHECK:           [[C0:%.*]] = constant 1 : index
+// CHECK:           [[C1:%.*]] = constant true
+// CHECK:           [[V0:%.*]] = scf.if [[C1]] -> (index) {
+// CHECK:             scf.yield [[C0]] : index
+// CHECK:           } else
+// CHECK:             scf.yield [[C0]] : index
+// CHECK:           }
+// CHECK:           return [[V0]] : index
+
+// -----
+
+func @nested_unused() -> (index) {
+  %c0 = constant 0 : index
+  %c1 = constant 1 : index
+  %true = constant true
+  %0, %1 = scf.if %true -> (index, index) {
+    %2, %3 = scf.if %true -> (index, index) {
+      scf.yield %c0, %c1 : index, index
+    } else {
+      scf.yield %c0, %c1 : index, index
+    }
+    scf.yield %2, %3 : index, index
+  } else {
+    scf.yield %c0, %c1 : index, index
+  }
+  return %1 : index
+}
+
+// CHECK-LABEL:   func @nested_unused
+// CHECK:           [[C0:%.*]] = constant 1 : index
+// CHECK:           [[C1:%.*]] = constant true
+// CHECK:           [[V0:%.*]] = scf.if [[C1]] -> (index) {
+// CHECK:             [[V1:%.*]] = scf.if [[C1]] -> (index) {
+// CHECK:               scf.yield [[C0]] : index
+// CHECK:             } else
+// CHECK:               scf.yield [[C0]] : index
+// CHECK:             }
+// CHECK:             scf.yield [[V1]] : index
+// CHECK:           } else
+// CHECK:             scf.yield [[C0]] : index
+// CHECK:           }
+// CHECK:           return [[V0]] : index
+
+// -----
+
+func @side_effect() {}
+func @all_unused() {
+  %c0 = constant 0 : index
+  %c1 = constant 1 : index
+  %true = constant true
+  %0, %1 = scf.if %true -> (index, index) {
+    call @side_effect() : () -> ()
+    scf.yield %c0, %c1 : index, index
+  } else {
+    call @side_effect() : () -> ()
+    scf.yield %c0, %c1 : index, index
+  }
+  return
+}
+
+// CHECK-LABEL:   func @all_unused
+// CHECK:           [[C1:%.*]] = constant true
+// CHECK:           scf.if [[C1]] {
+// CHECK:             call @side_effect() : () -> ()
+// CHECK:           } else
+// CHECK:             call @side_effect() : () -> ()
+// CHECK:           }
+// CHECK:           return