namespace {

#define GEN_PASS_DEF_TPUREWRITEPASS
#include "tensorflow/compiler/mlir/tensorflow/transforms/host_runtime/runtime_passes.h.inc"

struct TPURewritePass : public impl::TPURewritePassBase<TPURewritePass> {
  explicit TPURewritePass(llvm::StringRef _module_name)
      : module_name(_module_name) {}

  void runOnOperation() override;

  llvm::StringRef module_name;
};

    tf_device.launch() ops. This is because the embedding op is head outside
    compiled and the compile op is wrapped in launch to execute on host
    during TPURewritePass.

    For example, the tf.OpA with the `mini_batch_splits` attribute will be
    moved after _TPUCompileMlir and the first input will use the
    _TPUCompileMlir program output:

    ```mlir

  if (!read_var_op) return;

  // TODO(lyandy): Check if resource (first replica or replicate block arg)
  // shape has not changed in replicate prior to read. Currently after both
  // ResourceOpLiftingPass and TPURewritePass, there should not be any updates
  // to resources prior to their respective ReadVariableOp.
  if (auto block_arg =
          mlir::dyn_cast<BlockArgument>(read_var_op.getResource())) {

See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/

include "mlir/Pass/PassBase.td"

def TPURewritePass : Pass<"tf-tpu-rewrite", "mlir::ModuleOp"> {
  let summary = "Rewrites a `tf_device.cluster_func` on TPUs into TPU runtime operations.";

  let description = [{

 Both the embedding op and compile op are assumed to be wrapped in separate
 tf_device.launch() ops. This is because the embedding op is head outside
 compiled and the compile op is wrapped in launch to execute on host
 during TPURewritePass.

 For example, the tf.OpA with the `mini_batch_splits` attribute will be
 moved after _TPUCompileMlir and the first input will use the
 _TPUCompileMlir program output:

 ```mlir
 "tf_device.launch"() ({