namespace {

#define GEN_PASS_DEF_PARALLELEXECUTETOISLANDSPASS
#include "tensorflow/compiler/mlir/tensorflow/transforms/tf_passes.h.inc"

struct ParallelExecuteToIslandsPass
    : public impl::ParallelExecuteToIslandsPassBase<
          ParallelExecuteToIslandsPass> {
  explicit ParallelExecuteToIslandsPass(bool legacy_graph_export) {
    legacy_graph_export_ = legacy_graph_export;
  }
  void runOnOperation() override;
};

//    "tf.OpB"
//    "tf_device.parallel_execute"() ({
//      "tf.OpC"
//    }, {
//      "tf.OpD"
//    })
//  }
//
// The above IR will be flattened after `ReplicateToIslandPass` and
// `ParallelExecuteToIslandsPass` as follows:
//
// tf_executor.island wraps "tf.OpA"
// tf_executor.island {_parallel_execution_ids=r0:0} wraps "tf.OpB"
// tf_executor.island {_parallel_execution_ids=r0:0,p0:0} wraps "tf.OpC"

  let summary = "Fold explicit broadcasts into the following operations if they "
           "support implicit broadcasting on their operand.";
  let constructor = "TF::CreateBroadcastFoldPass()";
}

def ParallelExecuteToIslandsPass : Pass<"tf-parallel-execute-to-islands", "mlir::func::FuncOp"> {
  let summary = "Lowers device parallel_execute to executor islands";
  let constructor = "TFDevice::CreateParallelExecuteToIslandsPass()";