// RUN: tf-opt %s -split-input-file -verify-diagnostics -tf-embedding-pipelining | FILECHECK_OPTS="" FileCheck %s

// This test verifies the handling of TPU replicated inputs and outputs as well as the extraction of the four main functions.
module {
  func.func @main() {
    %cst_main = "tf.Const"() {value = dense<1> : tensor<i32>} : () -> tensor<i32>

  // Only one of EmbeddingSequencing and EmbeddingPipelining will actually
  // run and the logic is in EmbeddingPipeliningPass. If the pipelining pass
  // runs, embedding attributes are stripped and the sequencing pass will have
  // no effect. If the pipelining pass doesn't run, embedding attributes are
  // preserved and the sequencing rewrite will trigger.
  pm.addPass(mlir::TFDevice::CreateEmbeddingPipeliningPass());

// separate functions for proper sequencing of TF2 TPU Embedding (see
// tpu_embedding_v3.py). This pass is a precursor for pipelining (see
// embedding_pipelining.cc) and DOES NOT permit parallel execution across SC and
// TC. This pass is a temporary fallback to use while developing full pipelining
// capabilities.
//
// Ops are broken up into:
//   1. SC forward pass
//   2. TC forward/backward pass
//   3. SC backward pass

  // guarantee that tests are run on XLA and not on TF's CPU implementation.
  bool tf_xla_disable_constant_folding;

  // Disables full embedding pipelining when true. Instead, strict SparseCore
  // TensorCore sequencing will be used.
  bool tf_xla_disable_full_embedding_pipelining;

  // Force the WhileOps in embedding_pipelining and embedding_sequencing to use

            "compilation."),
       Flag("tf_xla_disable_full_embedding_pipelining",
            &build_ops_flags->tf_xla_disable_full_embedding_pipelining,
            "If true then disables full embedding pipelining and instead use "
            "strict SparseCore / TensorCore sequencing."),
       Flag("tf_xla_embedding_parallel_iterations",
            &build_ops_flags->tf_xla_embedding_parallel_iterations,

  // CHECK-NEXT:  return %3, %2 : tensor<3x3xf32>, tensor<*xf32>
  return %1#4, %1#3 : tensor<3x3xf32>, tensor<*xf32>
}

// -----
// This case test the re-ordering of the while body function to maximize pipelining between iterations.

// CHECK-LABEL: func.func private @"sort_map/while_cond"