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

using mlir::Operation;
using mlir::WalkResult;

class VerifyClusteringPass
    : public impl::VerifyClusteringPassBase<VerifyClusteringPass> {
 public:
  void runOnOperation() override;
};

void VerifyClusteringPass::runOnOperation() {
  Operation* func_op = getOperation();

  auto walk_result = func_op->walk([&](Operation* op) {

        ":tpu_sharding_identification_pass",
        ":verify_clustering_pass",
        ":xla_broadcast",
        ":xla_cluster_formation",
        "@llvm-project//mlir:FuncDialect",
        "@llvm-project//mlir:IR",
        "@llvm-project//mlir:Pass",
    ],
)

cc_library(
    name = "verify_clustering_pass",
    srcs = [
        "verify_clustering_pass.cc",
    ],
    deps = [
        ":clustering_passes_inc_gen",

// RUN: tf-opt -verify-clustering-pass  -split-input-file -verify-diagnostics %s | FileCheck %s
// Tests the VerifyClusteringPass Pass, ensures that an error is thrown when validation fails.

func.func @testNotTfDialect(%arg0: tensor<1x32x10x32xi32>, %arg1: tensor<32xi32>) -> tensor<1x32x10x32xi32> {
 // expected-error@below {{op is in dialect chlo not in tf functional dialect}}

// RUN: tf-opt -verify-input-dialect-to-executor-pass  -split-input-file -verify-diagnostics %s | FileCheck %s
// Tests the VerifyClusteringPass Pass, ensures that an error is thrown when validation fails.

// -----

// CHECK-LABEL: func @testNoClusterFuncOpPasses
func.func @testNoClusterFuncOpPasses(%arg0: tensor<4x?x!tf_type.stringref>) -> tensor<4x2x!tf_type.string> {
  %0 = "tf.Identity"(%arg0) : (tensor<4x?x!tf_type.stringref>) -> tensor<4x2x!tf_type.string>

See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
include "mlir/Pass/PassBase.td"

def VerifyClusteringPass : Pass<"verify-clustering-pass", "mlir::func::FuncOp"> {

  let summary = "Verify that the Bridge output is correct and errors if verification fails.";

  let description = [{