xref: /external/tensorflow/tensorflow/compiler/mlir/tfrt/transforms/merge_tf_if_ops.cc

/* Copyright 2021 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 "mlir/Dialect/StandardOps/IR/Ops.h"  // from @llvm-project
#include "mlir/Transforms/Passes.h"  // from @llvm-project
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
#include "tensorflow/compiler/mlir/tfrt/transforms/passes.h"

namespace tensorflow {
namespace tfrt_compiler {
namespace {

// A special DenseMapInfo that hashes only operands of a operation, and treats
// two operations equivalent if their operands are the same.
struct OpWithSameArgsInfo : llvm::DenseMapInfo<mlir::Operation *> {
  static unsigned getHashValue(const mlir::Operation *const_op) {
    auto *op = const_cast<mlir::Operation *>(const_op);
    return llvm::hash_combine(
        llvm::hash_combine_range(op->operand_begin(), op->operand_end()));
  }

  static bool isEqual(const mlir::Operation *const_lhs,
                      const mlir::Operation *const_rhs) {
    auto *lhs = const_cast<mlir::Operation *>(const_lhs);
    auto *rhs = const_cast<mlir::Operation *>(const_rhs);
    if (lhs == rhs) return true;
    if (lhs == getTombstoneKey() || lhs == getEmptyKey() ||
        rhs == getTombstoneKey() || rhs == getEmptyKey())
      return false;

    return std::equal(lhs->operand_begin(), lhs->operand_end(),
                      rhs->operand_begin(), rhs->operand_end());
  }
};

// This pass merges non-side-effecting tf.If ops if their operands are the same.
// For example,
//  %r0 = tf.If(%cond, %x) {else = @else_0, then = @then_0}
//  %r1, %r2 = tf.If(%cond, %x) {else = @else_1, then = @then_1}
//
// will be converted to:
//  func private @merge_else(%arg) {
//    %r0 = tf.PartitionedCall(%arg) {f = @else_0}
//    %r1, %r2 = tf.PartitionedCall(%arg) {f = @else_1}
//    return %r0, %r1, %r2
//  }
//  func private @merge_then(%arg) {
//    %r0 = tf.PartitionedCall(%arg) {f = @then_0}
//    %r1, %r2 = tf.PartitionedCall(%arg) {f = @then_1}
//    return %r0, %r1, %r2
//  }
//
//  %r0, %r1, %r2 = tf.If(%cond, %arg) {else = @merge_else, then = @merge_then}
//
// Note that the results will be concatenated.
class MergeTfIfOpsPass
    : public mlir::PassWrapper<MergeTfIfOpsPass,
                               mlir::OperationPass<mlir::ModuleOp>> {
  llvm::StringRef getArgument() const final { return "tfrt-merge-tf-if-ops"; }
  llvm::StringRef getDescription() const final {
    return "Merge stateless tf.If ops with the same arguments.";
  }

  void runOnOperation() override {
    auto module = getOperation();

    for (auto func_op :
         llvm::make_early_inc_range(module.getOps<mlir::FuncOp>())) {
      ProcessFunction(func_op);
    }
  }

  void ProcessFunction(mlir::FuncOp op) {
    // Use a hash map to group tf.If ops with the same operands.
    llvm::SmallDenseMap<mlir::Operation *, llvm::SmallVector<mlir::TF::IfOp, 2>,
                        2, OpWithSameArgsInfo>
        if_ops_to_merge;

    for (mlir::Operation &op : op.front()) {
      auto if_op = llvm::dyn_cast<mlir::TF::IfOp>(&op);

      // Skip non tf.If ops and tf.If ops that are side-effecting.
      if (!if_op || !if_op.is_stateless()) continue;

      if_ops_to_merge[if_op].push_back(if_op);
    }

    int id = 0;

    // Set the insertion point to the current function, as we will insert new
    // functions here.
    mlir::OpBuilder builder(op);

    // Track the tf.If ops that should be removed as they are merged.
    llvm::SmallVector<mlir::TF::IfOp, 4> if_ops_to_remove;

    for (auto &iter : if_ops_to_merge) {
      if (iter.second.size() <= 1) continue;

      // Merge tf.If ops that have the same operands. The merged branches will
      // be given unique names.
      MergeIfOpsWithSameArgs(
          builder, iter.first->getLoc(),
          /*branch_prefix=*/
          absl::StrCat(op.sym_name().str(), "_merged_if_", id++), iter.second);

      if_ops_to_remove.append(iter.second.begin(), iter.second.end());
    }

    // Now that we are no longer using `if_ops_to_merge` or any other data
    // structures that uses the operations that will be removed, we can now
    // erase these if ops.
    for (auto op : if_ops_to_remove) op->erase();
  }

  void MergeIfOpsWithSameArgs(mlir::OpBuilder &builder, mlir::Location loc,
                              absl::string_view branch_prefix,
                              llvm::MutableArrayRef<mlir::TF::IfOp> if_ops) {
    assert(if_ops.size() > 1);

    // The results of the merged tf.If op are the concatenation of results of
    // the original tf.If ops.
    llvm::SmallVector<mlir::Type, 4> new_result_types;
    for (auto if_op : if_ops) {
      new_result_types.append(if_op->result_type_begin(),
                              if_op->result_type_end());
    }

    auto branch_function_type = builder.getFunctionType(
        if_ops.front().input().getTypes(), new_result_types);

    // Create new branches for the merged tf.If op.
    auto then_branch_name = CreateBranchFunction(
        builder, loc, branch_prefix,
        /*branch_suffix=*/"_then", branch_function_type, if_ops,
        [](mlir::TF::IfOp op) { return op.then_branchAttr(); });

    auto else_branch_name = CreateBranchFunction(
        builder, loc, branch_prefix,
        /*branch_suffix=*/"_else", branch_function_type, if_ops,
        [](mlir::TF::IfOp op) { return op.else_branchAttr(); });

    mlir::OpBuilder::InsertionGuard guard(builder);
    builder.setInsertionPoint(if_ops.front());

    // Create the merged tf.If op using the new branches.
    auto new_if_op = builder.create<mlir::TF::IfOp>(
        loc, new_result_types, if_ops.front().cond(), if_ops.front().input(),
        then_branch_name, else_branch_name, /*is_stateless=*/true);

    // Replace the uses of results of the original tf.If ops with the results of
    // the merged tf.If op.
    auto new_result_iter = new_if_op.output().begin();
    for (auto if_op : if_ops) {
      for (auto result : if_op.output()) {
        assert(new_result_iter != new_if_op.output().end());
        result.replaceAllUsesWith(*new_result_iter);
        ++new_result_iter;
      }
    }
  }

  llvm::StringRef CreateBranchFunction(
      mlir::OpBuilder &builder, mlir::Location loc,
      absl::string_view branch_prefix, absl::string_view branch_suffix,
      mlir::FunctionType branch_function_type,
      llvm::ArrayRef<mlir::TF::IfOp> if_ops,
      llvm::function_ref<mlir::FlatSymbolRefAttr(mlir::TF::IfOp)> get_branch) {
    std::string branch_name = absl::StrCat(branch_prefix, branch_suffix);
    auto branch =
        builder.create<mlir::FuncOp>(loc, branch_name, branch_function_type);
    branch.setVisibility(mlir::FuncOp::Visibility::Private);

    mlir::OpBuilder::InsertionGuard guard(builder);

    // In the body of newly created branch function, we insert
    // tf.PartitionedCall ops to call the original branches.
    auto *block = branch.addEntryBlock();
    builder.setInsertionPointToStart(block);
    auto empty_string_attr = builder.getStringAttr("");

    llvm::SmallVector<mlir::Value, 4> results;
    results.reserve(branch_function_type.getNumResults());

    for (auto if_op : if_ops) {
      // Create the the call op to the original branch. The arguments are simply
      // the arguments from the wrapper function.
      auto call_op = builder.create<mlir::TF::PartitionedCallOp>(
          if_op.getLoc(), if_op.getResultTypes(), block->getArguments(),
          get_branch(if_op), empty_string_attr, empty_string_attr,
          empty_string_attr);

      // The results are the concatenation of the original branches.
      results.append(call_op.output().begin(), call_op.output().end());
    }

    builder.create<mlir::ReturnOp>(loc, results);

    return branch.sym_name();
  }
};

}  // namespace

std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>> CreateMergeTfIfOpsPass() {
  return std::make_unique<MergeTfIfOpsPass>();
}

static mlir::PassRegistration<MergeTfIfOpsPass> register_pass(
    CreateMergeTfIfOpsPass);

}  // namespace tfrt_compiler
}  // namespace tensorflow