// For architectures that support accelerated embedding lookups, this pass will
// rewrite the graph to use pipelining for better device utilization.
std::unique_ptr<OperationPass<ModuleOp>> CreateEmbeddingSequencingPass();

// This is a strictly sequential and formally correct fallback option for the
// embedding pipelining pass intended for debugging during pipelining
// development.
std::unique_ptr<OperationPass<ModuleOp>> CreateEmbeddingPipeliningPass();

include "mlir/Pass/PassBase.td"

def EmbeddingPipeliningPass : Pass<"tf-embedding-pipelining", "mlir::ModuleOp"> {
  let summary = "Rewrite graph for embedding pipelining";
  let constructor = "TFDevice::CreateEmbeddingPipeliningPass()";
    let description = [{
    For architectures that support accelerated embedding lookups, this pass will
    rewrite the graph to use pipelining for better device utilization.
  }];
}

// 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>

  // Enable automated pipelining pass unless:
  // 1. The user disables it via flag, or
  // 2. The graph contains TF.Summary ops. Graphs like this typically only run
  //    for a single step which doesn't work in pipelining.

  if (tensorflow::GetBuildXlaOpsPassFlags()
          ->tf_xla_disable_full_embedding_pipelining) {
    LOG(INFO) << "Embedding pipelining disabled via flag.";
    return false;
  }

  // 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());

          /*edge_filter=*/nullptr);
}

// This preserves the parallelism between pipeline stages.  For example, below
// is a typical pattern of input pipelining in Tensorflow and this heuristic
// ensures Node_X and Node_Y are put into different clusters.  Without the
// heuristic, they may be put into the same cluster and it can introduce

	}

	// TODO: only supports HTTP 1.1 for now.
	outReq.Proto = "HTTP/1.1"
	outReq.ProtoMajor = 1
	outReq.ProtoMinor = 1

	f.rewriter.Rewrite(outReq)

	// Disable closeNotify when method GET for http pipelining
	if outReq.Method == http.MethodGet {
		quietReq := outReq.WithContext(context.Background())
		*outReq = *quietReq
	}
}

// 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,