[mlir][SCF] Add utility to outline the then and else branches of an scf.IfOp

[jurahul]

mlir/lib/Dialect/SCF/Transforms/Utils.cpp

84

A utility to outline just a single Region will be useful by itself in other contexts. Can this be hoisted out to a common place (outside SCF dialect)?

[nicolasvasilache]

mlir/lib/Dialect/SCF/Transforms/Utils.cpp

84

Yes, but this will require handling multiple blocks which I will not have cycles to generalize before going on vacation.
At this point you may be better off refactoring / moving some code around once this lands if you need it?

[rriddle]

mlir/lib/Dialect/SCF/Transforms/Utils.cpp

109

This looks weird, can you just remap after construction or something? Cloning shouldn't be something considered this cheap.

[nicolasvasilache]

Drop extra clone.

[ftynse]

mlir/include/mlir/Dialect/SCF/Utils.h

57

Typo: "thepointer"

mlir/lib/Dialect/SCF/Transforms/Utils.cpp

90

Should the function rather take an OpBuilder as argument? This will make it compatible with the pattern rewriting infra.

96

What's the problem specifically? values.getTypes() should give you a TypeRange, and FunctionType::get accepts it.

106

Could this be simplified to for (Operation &op : ifOrElseRegion.front().without_terminator()) and a separate handling of the terminator?

111

If you lookup operands and put them in a vector, which is then passed into b.create, this will become compatible with pattern rewriting.

[nicolasvasilache]

Address review.

[nicolasvasilache]

DCE.

mlir/lib/Dialect/SCF/Transforms/Utils.cpp

96

Wouldn't convert automatically when this was originally written, works with rebase.

mlir/include/mlir/Dialect/SCF/Utils.h

 //===- Utils.h - SCF dialect utilities --------------------------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This header file defines prototypes for various SCF utilities.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_DIALECT_SCF_UTILS_H_
 #define MLIR_DIALECT_SCF_UTILS_H_
 
+#include "mlir/Support/LLVM.h"
+
 namespace mlir {
+class FuncOp;
 class OpBuilder;
 class ValueRange;
 
 namespace scf {
+class IfOp;
 class ForOp;
 class ParallelOp;
 } // end namespace scf
 
 /// Create a clone of `loop` with `newIterOperands` added as new initialization
 /// values and `newYieldedValues` added as new yielded values. The returned
 /// ForOp has `newYieldedValues.size()` new result values.  The `loop` induction
 /// variable and `newIterOperands` are remapped to the new induction variable

 /// legality checks.
 ///
 /// Prerequisite: `newYieldedValues.size() == newYieldedValues.size()`.
 scf::ForOp cloneWithNewYields(OpBuilder &b, scf::ForOp loop,
                               ValueRange newIterOperands,
                               ValueRange newYieldedValues,
                               bool replaceLoopResults = true);
 
+/// Outline the then and/or else regions of `ifOp` as follows:
+///  - if `thenFn` is not null, `thenFnName` must be specified and the `then`
+///    region is inlined into a new FuncOp that is captured by the pointer.
+///  - if `elseFn` is not null, `elseFnName` must be specified and the `else`
+///    region is inlined into a new FuncOp that is captured by the pointer.
+void outlineIfOp(OpBuilder &b, scf::IfOp ifOp, FuncOp *thenFn,
+                 StringRef thenFnName, FuncOp *elseFn, StringRef elseFnName);
 } // end namespace mlir
 #endif // MLIR_DIALECT_SCF_UTILS_H_

mlir/lib/Dialect/SCF/Transforms/CMakeLists.txt

add_mlir_dialect_library(MLIRSCFTransforms
 
   LINK_LIBS PUBLIC
   MLIRAffineOps
   MLIRIR
   MLIRPass
   MLIRSCF
   MLIRStandardOps
   MLIRSupport
-  )
+  MLIRTransformUtils

mlir/lib/Dialect/SCF/Transforms/Utils.cpp

 //===- LoopUtils.cpp ---- Misc utilities for loop transformation ----------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file implements miscellaneous loop transformation routines.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/SCF/Utils.h"
 
 #include "mlir/Dialect/SCF/SCF.h"
+#include "mlir/Dialect/StandardOps/IR/Ops.h"
 #include "mlir/IR/BlockAndValueMapping.h"
+#include "mlir/IR/Function.h"
+#include "mlir/Transforms/RegionUtils.h"
+
+#include "llvm/ADT/SetVector.h"
 
 using namespace mlir;
 
 scf::ForOp mlir::cloneWithNewYields(OpBuilder &b, scf::ForOp loop,
                                     ValueRange newIterOperands,
                                     ValueRange newYieldedValues,
                                     bool replaceLoopResults) {
   assert(newIterOperands.size() == newYieldedValues.size() &&
for (auto it : llvm::zip(loop.getResults(), newLoop.getResults().take_front(
       std::get<0>(it).replaceAllUsesWith(std::get<1>(it));
   }
 
   // TODO: this is unsafe in the context of a PatternRewrite.
   newYield.erase();
 
   return newLoop;
 }
+
+void mlir::outlineIfOp(OpBuilder &b, scf::IfOp ifOp, FuncOp *thenFn,
+                       StringRef thenFnName, FuncOp *elseFn,
+                       StringRef elseFnName) {
+  Location loc = ifOp.getLoc();
+  MLIRContext *ctx = ifOp.getContext();
+  auto outline = [&](Region &ifOrElseRegion, StringRef funcName) {
+    assert(!funcName.empty() && "Expected function name for outlining");
+    assert(ifOrElseRegion.getBlocks().size() <= 1 &&
+           "Expected at most one block");
+
+    // Outline before current function.
+    OpBuilder::InsertionGuard g(b);
+    b.setInsertionPoint(ifOp.getParentOfType<FuncOp>());
+
+    llvm::SetVector<Value> captures;
+    getUsedValuesDefinedAbove(ifOrElseRegion, captures);
+
+    ValueRange values(captures.getArrayRef());
+    FunctionType type =
+        FunctionType::get(values.getTypes(), ifOp.getResultTypes(), ctx);
+    auto outlinedFunc = b.create<FuncOp>(loc, funcName, type);
+    b.setInsertionPointToStart(outlinedFunc.addEntryBlock());
+    BlockAndValueMapping bvm;
+    for (auto it : llvm::zip(values, outlinedFunc.getArguments()))
+      bvm.map(std::get<0>(it), std::get<1>(it));
+    for (Operation &op : ifOrElseRegion.front().without_terminator())
+      b.clone(op, bvm);
+
+    Operation *term = ifOrElseRegion.front().getTerminator();
+    SmallVector<Value, 4> terminatorOperands;
+    for (auto op : term->getOperands())
+      terminatorOperands.push_back(bvm.lookup(op));
+    b.create<ReturnOp>(loc, term->getResultTypes(), terminatorOperands);
+
+    ifOrElseRegion.front().clear();
+    b.setInsertionPointToEnd(&ifOrElseRegion.front());
+    Operation *call = b.create<CallOp>(loc, outlinedFunc, values);
+    b.create<scf::YieldOp>(loc, call->getResults());
+    return outlinedFunc;
+  };
+
+  if (thenFn && !ifOp.thenRegion().empty())
+    *thenFn = outline(ifOp.thenRegion(), thenFnName);
+  if (elseFn && !ifOp.elseRegion().empty())
+    *elseFn = outline(ifOp.elseRegion(), elseFnName);
+}

mlir/test/Transforms/scf-if-utils.mlir

+// RUN: mlir-opt -allow-unregistered-dialect -test-scf-if-utils -split-input-file %s | FileCheck %s
+
+// -----
+
+//      CHECK: func @outlined_then0(%{{.*}}: i1, %{{.*}}: memref<?xf32>) -> i8 {
+// CHECK-NEXT:   %{{.*}} = "some_op"(%{{.*}}, %{{.*}}) : (i1, memref<?xf32>) -> i8
+// CHECK-NEXT:   return %{{.*}} : i8
+// CHECK-NEXT: }
+//      CHECK: func @outlined_else0(%{{.*}}: i8) -> i8 {
+// CHECK-NEXT:   return %{{.*}}0 : i8
+// CHECK-NEXT: }
+//      CHECK: func @outline_if_else(
+// CHECK-NEXT:   %{{.*}} = scf.if %{{.*}} -> (i8) {
+// CHECK-NEXT:     %{{.*}} = call @outlined_then0(%{{.*}}, %{{.*}}) : (i1, memref<?xf32>) -> i8
+// CHECK-NEXT:     scf.yield %{{.*}} : i8
+// CHECK-NEXT:   } else {
+// CHECK-NEXT:     %{{.*}} = call @outlined_else0(%{{.*}}) : (i8) -> i8
+// CHECK-NEXT:     scf.yield %{{.*}} : i8
+// CHECK-NEXT:   }
+// CHECK-NEXT:   return
+// CHECK-NEXT: }
+func @outline_if_else(%cond: i1, %a: index, %b: memref<?xf32>, %c: i8) {
+  %r = scf.if %cond -> (i8) {
+    %r = "some_op"(%cond, %b) : (i1, memref<?xf32>) -> (i8)
+    scf.yield %r : i8
+  } else {
+    scf.yield %c : i8
+  }
+  return
+}
+
+// -----
+
+//      CHECK: func @outlined_then0(%{{.*}}: i1, %{{.*}}: memref<?xf32>) {
+// CHECK-NEXT:   "some_op"(%{{.*}}, %{{.*}}) : (i1, memref<?xf32>) -> ()
+// CHECK-NEXT:   return
+// CHECK-NEXT: }
+//      CHECK: func @outline_if(
+// CHECK-NEXT:   scf.if %{{.*}} {
+// CHECK-NEXT:     call @outlined_then0(%{{.*}}, %{{.*}}) : (i1, memref<?xf32>) -> ()
+// CHECK-NEXT:   }
+// CHECK-NEXT:   return
+// CHECK-NEXT: }
+func @outline_if(%cond: i1, %a: index, %b: memref<?xf32>, %c: i8) {
+  scf.if %cond {
+    "some_op"(%cond, %b) : (i1, memref<?xf32>) -> ()
+    scf.yield
+  }
+  return
+}
+
+// -----
+
+//      CHECK: func @outlined_then0() {
+// CHECK-NEXT:   return
+// CHECK-NEXT: }
+//      CHECK: func @outlined_else0(%{{.*}}: i1, %{{.*}}: memref<?xf32>) {
+// CHECK-NEXT:   "some_op"(%{{.*}}, %{{.*}}) : (i1, memref<?xf32>) -> ()
+// CHECK-NEXT:   return
+// CHECK-NEXT: }
+//      CHECK: func @outline_empty_if_else(
+// CHECK-NEXT:   scf.if %{{.*}} {
+// CHECK-NEXT:     call @outlined_then0() : () -> ()
+// CHECK-NEXT:   } else {
+// CHECK-NEXT:     call @outlined_else0(%{{.*}}, %{{.*}}) : (i1, memref<?xf32>) -> ()
+// CHECK-NEXT:   }
+// CHECK-NEXT:   return
+// CHECK-NEXT: }
+func @outline_empty_if_else(%cond: i1, %a: index, %b: memref<?xf32>, %c: i8) {
+  scf.if %cond {
+  } else {
+    "some_op"(%cond, %b) : (i1, memref<?xf32>) -> ()
+  }
+  return
+}

mlir/test/Transforms/scf-loop-utils.mlir

-// RUN: mlir-opt -allow-unregistered-dialect -test-scf-utils -mlir-disable-threading %s | FileCheck %s
+// RUN: mlir-opt -allow-unregistered-dialect -test-scf-for-utils -mlir-disable-threading %s | FileCheck %s
 
 // CHECK-LABEL: @hoist
 //  CHECK-SAME: %[[lb:[a-zA-Z0-9]*]]: index,
 //  CHECK-SAME: %[[ub:[a-zA-Z0-9]*]]: index,
 //  CHECK-SAME: %[[step:[a-zA-Z0-9]*]]: index
 func @hoist(%lb: index, %ub: index, %step: index) {
   // CHECK: %[[A:.*]] = "fake_read"() : () -> index
   // CHECK: %[[RES:.*]] = scf.for %[[I:.*]] = %[[lb]] to %[[ub]] step %[[step]] iter_args(%[[VAL:.*]] = %[[A]]) -> (index)

mlir/test/lib/Transforms/TestSCFUtils.cpp

 #include "mlir/Pass/Pass.h"
 #include "mlir/Transforms/Passes.h"
 
 #include "llvm/ADT/SetVector.h"
 
 using namespace mlir;
 
 namespace {
-class TestSCFUtilsPass : public PassWrapper<TestSCFUtilsPass, FunctionPass> {
+class TestSCFForUtilsPass
+    : public PassWrapper<TestSCFForUtilsPass, FunctionPass> {
 public:
-  explicit TestSCFUtilsPass() {}
+  explicit TestSCFForUtilsPass() {}
 
   void runOnFunction() override {
     FuncOp func = getFunction();
     SmallVector<scf::ForOp, 4> toErase;
 
     func.walk([&](Operation *fakeRead) {
       if (fakeRead->getName().getStringRef() != "fake_read")
         return;
func.walk([&](Operation *fakeRead) {
       fakeCompute->getResult(0).replaceAllUsesWith(
           newLoop.getResults().take_back()[0]);
       toErase.push_back(loop);
     });
     for (auto loop : llvm::reverse(toErase))
       loop.erase();
   }
 };
+
+class TestSCFIfUtilsPass
+    : public PassWrapper<TestSCFIfUtilsPass, FunctionPass> {
+public:
+  explicit TestSCFIfUtilsPass() {}
+
+  void runOnFunction() override {
+    int count = 0;
+    FuncOp func = getFunction();
+    func.walk([&](scf::IfOp ifOp) {
+      auto strCount = std::to_string(count++);
+      FuncOp thenFn, elseFn;
+      OpBuilder b(ifOp);
+      outlineIfOp(b, ifOp, &thenFn, std::string("outlined_then") + strCount,
+                  &elseFn, std::string("outlined_else") + strCount);
+    });
+  }
+};
 } // end namespace
 
 namespace mlir {
 void registerTestSCFUtilsPass() {
-  PassRegistration<TestSCFUtilsPass>("test-scf-utils", "test scf utils");
+  PassRegistration<TestSCFForUtilsPass>("test-scf-for-utils",
+                                        "test scf.for utils");
+  PassRegistration<TestSCFIfUtilsPass>("test-scf-if-utils",
+                                       "test scf.if utils");
 }
 } // namespace mlir