Mason Chang <******@****.***> 1714440931 -0700

        "//tensorflow/compiler/mlir/tf2xla/internal:__subpackages__",
    ],
    licenses = ["notice"],
)

cc_library(
    name = "clustering_passes",
    hdrs = [
        "clustering_passes.h",
    ],
    textual_hdrs = [
        "clustering_passes.h.inc",
    ],
    deps = [
        ":extract_head_tail_outside_compilation",
        ":extract_outside_compilation",
        ":hoist_broadcast_read",

#include "tensorflow/compiler/mlir/tensorflow/transforms/passes.h"
#include "tensorflow/compiler/mlir/tensorflow/transforms/sparsecore/sparsecore_passes.h"
#include "tensorflow/compiler/mlir/tf2xla/internal/passes/clustering_passes.h"

namespace tensorflow {
namespace tf2xla {
namespace internal {

using mlir::OpPassManager;
using mlir::func::FuncOp;

// Adds replicated Bridge clustering pipeline passes to the given pass_manager.

using mlir::tf_device::LaunchOp;
using mlir::tf_device::ParallelExecuteOp;
using mlir::tf_device::ReplicateOp;

#define GEN_PASS_DEF_XLABROADCASTPASS
#include "tensorflow/compiler/mlir/tf2xla/internal/passes/clustering_passes.h.inc"

const char kICIWeightDistributionMlirBridgeMarker[] =
    "ici_weight_distribution_mlir_bridge_marker";

struct XlaBroadcast : public impl::XlaBroadcastPassBase<XlaBroadcast> {

    std::get<0>(result).replaceAllUsesWith(std::get<1>(result));

  cluster.erase();
}

#define GEN_PASS_DEF_EXTRACTHEADTAILOUTSIDECOMPILATIONPASS
#include "tensorflow/compiler/mlir/tf2xla/internal/passes/clustering_passes.h.inc"

struct ExtractHeadTailOutsideCompilationPass
    : public impl::ExtractHeadTailOutsideCompilationPassBase<
          ExtractHeadTailOutsideCompilationPass> {
  void runOnOperation() override;
};

        "Tracks if auto outside compilation is enabled");

#define GEN_PASS_DEF_MARKOPSFOROUTSIDECOMPILATIONPASS
#include "tensorflow/compiler/mlir/tf2xla/internal/passes/clustering_passes.h.inc"

struct MarkOpsForOutsideCompilation
    : public impl::MarkOpsForOutsideCompilationPassBase<
          MarkOpsForOutsideCompilation> {
  void runOnOperation() override;
};

constexpr char kAliasingAttr[] = "tf.aliasing_output";
constexpr char kNumCoresPerReplicaAttr[] = "num_cores_per_replica";

#define GEN_PASS_DEF_TPUSHARDINGIDENTIFICATIONPASS
#include "tensorflow/compiler/mlir/tf2xla/internal/passes/clustering_passes.h.inc"

struct TPUShardingIdentificationPass
    : public impl::TPUShardingIdentificationPassBase<
          TPUShardingIdentificationPass> {
  void runOnOperation() final;
};

using ClusterMap = llvm::SmallDenseMap<llvm::StringRef, OpSetVector, 8>;

#define GEN_PASS_DEF_TPUCLUSTERFORMATIONPASS
#include "tensorflow/compiler/mlir/tf2xla/internal/passes/clustering_passes.h.inc"

class TPUClusterFormationPass
    : public impl::TPUClusterFormationPassBase<TPUClusterFormationPass> {
 public:
  explicit TPUClusterFormationPass(bool strict_clusters)