xref: /external/tensorflow/tensorflow/compiler/mlir/tensorflow/transforms/merge_control_flow.cc

/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/

#include <memory>
#include <queue>
#include <string>
#include <utility>

#include "llvm/Support/FormatVariadic.h"
#include "mlir/IR/TypeUtilities.h"  // from @llvm-project
#include "mlir/Pass/Pass.h"  // from @llvm-project
#include "mlir/Pass/PassRegistry.h"  // from @llvm-project
#include "mlir/Transforms/RegionUtils.h"  // from @llvm-project
#include "tensorflow/compiler/mlir/tensorflow/analysis/side_effect_analysis.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_device.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_types.h"
#include "tensorflow/compiler/mlir/tensorflow/transforms/passes.h"
#include "tensorflow/compiler/mlir/xla/transforms/passes.h"

namespace mlir {
namespace TFDevice {

namespace {

// This pass merges IfRegion ops together if they have the same predicate and it
// is safe to do so (there are no intermediate dependencies, they are in the
// same block, etc).
//
// A simple example:
//    "tf.IfRegion"(%0) ( {
//      %2 = "tf.A"() : () -> (tensor<f32>)
//      "tf.Yield"() : () -> ()
//      }, {
//      "tf.Yield"() : () -> ()
//     }) { is_stateless = true } : (tensor<i1>) -> ()
//    "tf.IfRegion"(%0) ( {
//      %2 = "tf.B"() : () -> (tensor<f32>)
//      "tf.Yield"() : () -> ()
//      }, {
//      "tf.Yield"() : () -> ()
//     }) { is_stateless = true } : (tensor<i1>) -> ()
// Would become:
//    "tf.IfRegion"(%0) ( {
//      %2 = "tf.A"() : () -> (tensor<f32>)
//      %3 = "tf.B"() : () -> (tensor<f32>)
//      "tf.Yield"() : () -> ()
//      }, {
//      "tf.Yield"() : () -> ()
//     }) { is_stateless = true } : (tensor<i1>) -> ()

struct MergeControlFlow : public TF::PerFunctionAggregateAnalysisConsumerPass<
                              MergeControlFlow, TF::SideEffectAnalysis> {
  void runOnFunction(FuncOp func,
                     const TF::SideEffectAnalysis::Info& side_effect_analysis);
};

// Returns whether it is safe to merge `source` IfRegion into `destination`
// IfRegion. `source` must come after `destination`.
bool SafeToMerge(TF::IfRegionOp source, TF::IfRegionOp destination,
                 const TF::SideEffectAnalysis::Info& side_effect_analysis) {
  // IfRegion ops must be in the same block.
  if (source.getOperation()->getBlock() !=
      destination.getOperation()->getBlock())
    return false;
  assert(destination.getOperation()->isBeforeInBlock(source.getOperation()));

  llvm::SmallSetVector<Operation*, 4> source_ops;
  source_ops.insert(source);
  for (Operation& op : source.then_branch().front()) {
    source_ops.insert(&op);
  }
  for (Operation& op : source.else_branch().front()) {
    source_ops.insert(&op);
  }

  // If there is an intermediate data or side effect dependency between the
  // ops in destination and the ops in the source, it's not safe to merge
  // them.
  std::vector<Operation*> dependencies;
  for (auto* user : destination.getOperation()->getUsers()) {
    if (!source_ops.contains(user)) dependencies.push_back(user);
  }
  for (Operation& op : destination.then_branch().front()) {
    for (auto* successor : side_effect_analysis.DirectControlSuccessors(&op)) {
      if (!source_ops.contains(successor)) dependencies.push_back(successor);
    }
  }
  for (Operation& op : destination.else_branch().front()) {
    for (auto* successor : side_effect_analysis.DirectControlSuccessors(&op)) {
      if (!source_ops.contains(successor)) dependencies.push_back(successor);
    }
  }

  bool safe_to_merge = true;

  llvm::SmallPtrSet<Operation*, 4> visited;
  while (!dependencies.empty()) {
    Operation* dependency = dependencies.back();
    dependencies.pop_back();
    if (visited.count(dependency)) continue;
    visited.insert(dependency);
    for (auto* user : dependency->getUsers()) {
      if (source_ops.contains(user)) {
        safe_to_merge = false;
        break;
      } else {
        dependencies.push_back(user);
      }
    }
    for (auto* successor :
         side_effect_analysis.DirectControlSuccessors(dependency)) {
      if (source_ops.contains(successor)) {
        safe_to_merge = false;
        break;
      } else {
        dependencies.push_back(successor);
      }
    }
    // If the op is nested, then also consider the users and successors of the
    // parent op.
    if (dependency->getBlock() != destination.getOperation()->getBlock())
      dependencies.push_back(dependency->getParentOp());
    if (!safe_to_merge) break;
  }
  return safe_to_merge;
}

// Checks whether a return indice should be kep for `first_if_op` by checking
// for results in `second_if_op`.
llvm::SmallVector<int, 4> GetReturnIndicesToKeep(TF::IfRegionOp first_if_op,
                                                 TF::IfRegionOp second_if_op) {
  llvm::SmallVector<int, 4> return_indices_to_keep;
  for (auto& index_and_value : llvm::enumerate(first_if_op.getResults())) {
    if (!llvm::all_of(index_and_value.value().getUsers(), [&](Operation* op) {
          return second_if_op->isProperAncestor(op);
        })) {
      return_indices_to_keep.push_back(index_and_value.index());
    }
  }
  return return_indices_to_keep;
}

// Move the body excluding the terminators of else and then regions from
// 'source' to 'destination'.
void MoveBranches(TF::IfRegionOp source, TF::IfRegionOp destination) {
  Block& destination_then_block = destination.then_branch().front();
  auto& source_then_body = source.then_branch().front().getOperations();
  destination_then_block.getOperations().splice(
      destination_then_block.without_terminator().end(), source_then_body,
      source_then_body.begin(), std::prev(source_then_body.end()));

  Block& destination_else_block = destination.else_branch().front();
  auto& source_else_body = source.else_branch().front().getOperations();
  destination_else_block.getOperations().splice(
      destination_else_block.without_terminator().end(), source_else_body,
      source_else_body.begin(), std::prev(source_else_body.end()));
}

// Move all ops that depends on the results from `result_op` after `after_op`.
void MoveResultsAfter(
    Operation* result_op, Operation* after_op,
    const TF::SideEffectAnalysis::Info& side_effect_analysis) {
  std::queue<Operation*> queue;

  auto enqueue_deps = [&](Operation* source_op) {
    for (Operation* user : source_op->getUsers()) {
      queue.push(user);
    }
    source_op->walk([&](Operation* walked_op) {
      for (Operation* successor :
           side_effect_analysis.DirectControlSuccessors(walked_op)) {
        if (!source_op->isProperAncestor(successor)) queue.push(successor);
      }
    });
  };
  enqueue_deps(result_op);

  while (!queue.empty()) {
    auto* op = queue.front();
    queue.pop();
    while (op->getBlock() != after_op->getBlock()) op = op->getParentOp();
    if (op->isBeforeInBlock(after_op)) {
      op->moveAfter(after_op);
      after_op = op;
      enqueue_deps(op);
    }
  }
}

TF::IfRegionOp CreateMergedIf(
    ArrayRef<int> source_return_indices_to_keep,
    ArrayRef<int> destination_return_indices_to_keep, TF::IfRegionOp source,
    TF::IfRegionOp destination,
    const TF::SideEffectAnalysis::Info& side_effect_analysis) {
  llvm::SmallVector<Type, 4> merged_return_types;
  for (int i : destination_return_indices_to_keep)
    merged_return_types.push_back(destination.getResult(i).getType());
  for (int i : source_return_indices_to_keep)
    merged_return_types.push_back(source.getResult(i).getType());

  OpBuilder builder(destination);
  // Create new IfRegion with correct merged results.
  builder.setInsertionPoint(source.getOperation());

  auto new_if_op = builder.create<TF::IfRegionOp>(
      destination.getLoc(), merged_return_types, destination.cond(),
      destination.is_stateless() && source.is_stateless(),
      destination._then_func_nameAttr(), destination._else_func_nameAttr());
  new_if_op.then_branch().push_back(new Block);
  new_if_op.else_branch().push_back(new Block);
  // Replace internal usages of merged if ops.
  for (OpResult result : destination.getResults()) {
    replaceAllUsesInRegionWith(
        result,
        destination.then_branch().front().getTerminator()->getOperand(
            result.getResultNumber()),
        source.then_branch());
    replaceAllUsesInRegionWith(
        result,
        destination.else_branch().front().getTerminator()->getOperand(
            result.getResultNumber()),
        source.else_branch());
  }

  MoveResultsAfter(destination.getOperation(), new_if_op.getOperation(),
                   side_effect_analysis);

  // Replace external usages of merged if ops.
  int new_return_index = 0;
  for (int i : destination_return_indices_to_keep) {
    destination.getResult(i).replaceAllUsesWith(
        new_if_op.getResult(new_return_index++));
  }
  for (int i : source_return_indices_to_keep) {
    source.getResult(i).replaceAllUsesWith(
        new_if_op.getResult(new_return_index++));
  }

  // Create the Yield ops for both branches with merged results.
  llvm::SmallVector<Value, 4> merged_then_yield_values;
  for (int i : destination_return_indices_to_keep)
    merged_then_yield_values.push_back(
        destination.then_branch().front().getTerminator()->getOperand(i));
  for (int i : source_return_indices_to_keep)
    merged_then_yield_values.push_back(
        source.then_branch().front().getTerminator()->getOperand(i));
  builder.setInsertionPointToEnd(&new_if_op.then_branch().front());
  builder.create<TF::YieldOp>(
      destination.then_branch().front().getTerminator()->getLoc(),
      /*operands=*/merged_then_yield_values);

  llvm::SmallVector<Value, 4> merged_else_yield_values;
  for (int i : destination_return_indices_to_keep)
    merged_else_yield_values.push_back(
        destination.else_branch().front().getTerminator()->getOperand(i));
  for (int i : source_return_indices_to_keep)
    merged_else_yield_values.push_back(
        source.else_branch().front().getTerminator()->getOperand(i));
  builder.setInsertionPointToEnd(&new_if_op.else_branch().front());
  builder.create<TF::YieldOp>(
      destination.else_branch().front().getTerminator()->getLoc(),
      /*operands=*/merged_else_yield_values);

  // Merge the two branch regions from both IfRegionOps into new IfRegionOp.
  MoveBranches(/*source=*/destination, /*destination=*/new_if_op);
  destination.erase();
  MoveBranches(/*source=*/source, /*destination=*/new_if_op);
  source.erase();
  return new_if_op;
}

// Groups if regions by common predicate and attemps to merge them.
void OptimizeIfRegions(
    Block* block, const TF::SideEffectAnalysis::Info& side_effect_analysis) {
  // Determine IfRegions with the same predicate.
  llvm::SmallDenseMap<Value, llvm::SmallVector<TF::IfRegionOp, 8>, 8>
      grouped_if_ops;
  block->walk([&](TF::IfRegionOp if_op) {
    auto it = grouped_if_ops.try_emplace(if_op.cond());
    it.first->getSecond().push_back(if_op);
  });

  for (auto& entry : grouped_if_ops) {
    auto& if_ops = entry.second;
    for (auto it = if_ops.begin(); it != if_ops.end(); ++it) {
      TF::IfRegionOp first_if_op = *it;
      for (auto it2 = std::next(it); it2 != if_ops.end(); ++it2) {
        TF::IfRegionOp second_if_op = *it2;
        if (!SafeToMerge(second_if_op, first_if_op, side_effect_analysis))
          break;

        // For both check if there are uses outside of IfRegion, keep these as
        // part of the return and replace the internal uses.
        auto first_return_indices_to_keep =
            GetReturnIndicesToKeep(first_if_op, second_if_op);
        auto second_return_indices_to_keep =
            GetReturnIndicesToKeep(second_if_op, first_if_op);

        auto new_if_op = CreateMergedIf(
            second_return_indices_to_keep, first_return_indices_to_keep,
            second_if_op, first_if_op, side_effect_analysis);

        if_ops.erase(it2--);
        first_if_op = new_if_op;
      }
    }
  }
}

void MergeControlFlow::runOnFunction(
    FuncOp func, const TF::SideEffectAnalysis::Info& side_effect_analysis) {
  auto result = func.walk([&](tf_device::ClusterOp cluster) {
    OptimizeIfRegions(&cluster.GetBody(), side_effect_analysis);
    return WalkResult::advance();
  });

  if (result.wasInterrupted()) return signalPassFailure();
}

}  // namespace

std::unique_ptr<OperationPass<ModuleOp>> CreateMergeControlFlowPass() {
  return std::make_unique<MergeControlFlow>();
}

static PassRegistration<MergeControlFlow> pass(
    "tf-merge-control-flow", "Merges control flow with a common predicate.");
}  // namespace TFDevice
}  // namespace mlir