namespace mlir {
namespace TFDevice {

namespace {

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

struct ClusterFormationPass
    : public impl::ClusterFormationPassBase<ClusterFormationPass> {
  void runOnOperation() override;
};

void ReplaceLiveOutExternalUses(llvm::ArrayRef<Value> live_outs,
                                tf_device::LaunchOp launch_op) {

  let summary = "Unroll TF BatchMatMul op into Reshape, Slice, MatMul, Pack ops.";
  let constructor = "TF::CreateUnrollBatchMatMulPassPass()";
}

def ClusterFormationPass : Pass<"tf-device-cluster-formation", "mlir::ModuleOp"> {
  let summary = "Form clusters from instructions assigned to same device";
  let constructor = "TFDevice::CreateClusterFormationPass()";